59 lines
2.5 KiB
Python
59 lines
2.5 KiB
Python
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, spectral_norm=True):
|
|
padding = (kernel_size // 2) * dilation
|
|
conv_layer = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding)
|
|
if spectral_norm:
|
|
conv_layer = utils.spectral_norm(conv_layer)
|
|
return nn.Sequential(
|
|
conv_layer,
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
nn.BatchNorm1d(out_channels)
|
|
)
|
|
|
|
class AttentionBlock(nn.Module):
|
|
def __init__(self, channels):
|
|
super(AttentionBlock, self).__init__()
|
|
self.attention = nn.Sequential(
|
|
nn.Conv1d(channels, channels // 4, kernel_size=1),
|
|
nn.ReLU(),
|
|
nn.Conv1d(channels // 4, channels, kernel_size=1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x):
|
|
attention_weights = self.attention(x)
|
|
return x * attention_weights
|
|
|
|
class SISUDiscriminator(nn.Module):
|
|
def __init__(self, layers=64): #Increased base layer count
|
|
super(SISUDiscriminator, self).__init__()
|
|
self.model = nn.Sequential(
|
|
discriminator_block(1, layers, kernel_size=3, stride=1), #Aggressive downsampling
|
|
discriminator_block(layers, layers * 2, kernel_size=5, stride=2),
|
|
discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=4),
|
|
|
|
#AttentionBlock(layers * 4), #Added attention
|
|
|
|
#discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=4),
|
|
#AttentionBlock(layers * 8), #Added attention
|
|
#discriminator_block(layers * 8, layers * 16, kernel_size=5, dilation=8),
|
|
#discriminator_block(layers * 16, layers * 16, kernel_size=3, dilation=1),
|
|
#discriminator_block(layers * 16, layers * 8, kernel_size=3, dilation=2),
|
|
#discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=1),
|
|
discriminator_block(layers * 4, layers * 2, kernel_size=5, stride=2),
|
|
discriminator_block(layers * 2, layers, kernel_size=3, stride=1),
|
|
discriminator_block(layers, 1, kernel_size=3, stride=1, spectral_norm=False) #last layer no spectral norm.
|
|
)
|
|
self.global_avg_pool = nn.AdaptiveAvgPool1d(1)
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
x = self.model(x)
|
|
x = self.global_avg_pool(x)
|
|
x = x.view(-1, 1)
|
|
x = self.sigmoid(x)
|
|
return x
|