import torch import torch.nn as nn import torch.nn.utils as utils def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dilation=1): padding = (kernel_size // 2) * dilation return nn.Sequential( utils.spectral_norm(nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding)), nn.LeakyReLU(0.2, inplace=True), nn.BatchNorm1d(out_channels) ) class SISUDiscriminator(nn.Module): def __init__(self): super(SISUDiscriminator, self).__init__() layers = 4 # Increased base layer count self.model = nn.Sequential( discriminator_block(1, layers, kernel_size=7, stride=2), # Initial downsampling discriminator_block(layers, layers * 2, kernel_size=5, stride=2), # Downsampling discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=2), # Increased dilation discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=4), # Increased dilation discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=8), # Deeper layer! discriminator_block(layers * 8, layers * 8, kernel_size=5, dilation=1), # Deeper layer! discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=2), # Reduced dilation discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=1), discriminator_block(layers * 2, layers, kernel_size=3, stride=1), # Final convolution discriminator_block(layers, 1, kernel_size=3, stride=1) ) self.global_avg_pool = nn.AdaptiveAvgPool1d(1) def forward(self, x): # Gaussian noise is not necessary here for discriminator as it is already implicit in the training process x = self.model(x) x = self.global_avg_pool(x) x = x.view(-1, 1) return x