33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import torch.nn as nn
|
|
|
|
class SISUGenerator(nn.Module):
|
|
def __init__(self, upscale_scale=4): # No noise_dim parameter
|
|
super(SISUGenerator, self).__init__()
|
|
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()
|
|
)
|
|
|
|
# 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)
|
|
)
|
|
|
|
def forward(self, x, upscale_scale=4):
|
|
x = self.conv1(x)
|
|
x = self.upsample(x)
|
|
x = self.conv2(x)
|
|
return x
|