NikkeDoy/SISU
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

:albemic: | Experimenting with other model layouts.

Browse Source
This commit is contained in:
NikkeDoy
2025-03-15 18:01:19 +02:00
parent 416500f7fc
commit 7e1c7e935a
4 changed files with 96 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
discriminator.py
View File

@ -2,35 +2,54 @@ 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):
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(
utils.spectral_norm(nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding)),
conv_layer,
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)
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()
)
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
attention_weights = self.attention(x)
return x * attention_weights
class SISUDiscriminator(nn.Module):
def __init__(self, layers=4): #Increased base layer count
super(SISUDiscriminator, self).__init__()
self.model = nn.Sequential(
discriminator_block(1, layers, kernel_size=7, stride=4), #Aggressive downsampling
discriminator_block(layers, layers * 2, kernel_size=5, stride=2),
discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=2),
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=3, stride=1),
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