⚗️ | Experimenting still...
This commit is contained in:
53
generator.py
53
generator.py
@@ -1,39 +1,32 @@
|
||||
import torch.nn as nn
|
||||
|
||||
class SISUGenerator(nn.Module):
|
||||
def __init__(self, upscale_scale=1):
|
||||
def __init__(self, upscale_scale=4): # No noise_dim parameter
|
||||
super(SISUGenerator, self).__init__()
|
||||
self.layers1 = nn.Sequential(
|
||||
nn.Conv1d(2, 128, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True), # Activation
|
||||
nn.BatchNorm1d(128), # Batch Norm
|
||||
nn.Conv1d(128, 256, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True), # Activation
|
||||
nn.BatchNorm1d(256), # Batch Norm
|
||||
layer = 32
|
||||
# Convolution layers
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv1d(1, layer * 2, kernel_size=7, padding=1),
|
||||
nn.PReLU(),
|
||||
nn.Conv1d(layer * 2, layer * 5, kernel_size=5, padding=1),
|
||||
nn.PReLU(),
|
||||
nn.Conv1d(layer * 5, layer * 5, kernel_size=3, padding=1),
|
||||
nn.PReLU()
|
||||
)
|
||||
|
||||
self.layers2 = nn.Sequential(
|
||||
nn.Conv1d(256, 128, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True), # Activation
|
||||
nn.BatchNorm1d(128), # Batch Norm
|
||||
nn.Conv1d(128, 64, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True), # Activation
|
||||
nn.BatchNorm1d(64), # Batch Norm
|
||||
nn.Conv1d(64, upscale_scale * 2, kernel_size=3, padding=1), # Output channels scaled
|
||||
# Transposed convolution for upsampling
|
||||
self.upsample = nn.ConvTranspose1d(layer * 5, layer * 5, kernel_size=upscale_scale, stride=upscale_scale)
|
||||
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv1d(layer * 5, layer * 5, kernel_size=3, padding=1),
|
||||
nn.PReLU(),
|
||||
nn.Conv1d(layer * 5, layer * 2, kernel_size=5, padding=1),
|
||||
nn.PReLU(),
|
||||
nn.Conv1d(layer * 2, 1, kernel_size=7, padding=1)
|
||||
)
|
||||
self.upscale_factor = upscale_scale
|
||||
|
||||
def pixel_shuffle_1d(self, input, upscale_factor):
|
||||
batch_size, channels, in_width = input.size()
|
||||
out_width = in_width * upscale_factor
|
||||
input_view = input.contiguous().view(batch_size, channels // upscale_factor, upscale_factor, in_width)
|
||||
shuffle_out = input_view.permute(0, 1, 3, 2).contiguous()
|
||||
return shuffle_out.view(batch_size, channels // upscale_factor, out_width)
|
||||
|
||||
def forward(self, x, scale):
|
||||
x = self.layers1(x)
|
||||
upsample = nn.Upsample(scale_factor=scale, mode='nearest')
|
||||
x = upsample(x)
|
||||
x = self.layers2(x)
|
||||
x = self.pixel_shuffle_1d(x, self.upscale_factor)
|
||||
def forward(self, x, upscale_scale=4):
|
||||
x = self.conv1(x)
|
||||
x = self.upsample(x)
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
Reference in New Issue
Block a user