:albemic: | Tests.

This commit is contained in:
2025-03-25 19:50:51 +02:00
parent 7e1c7e935a
commit 54338e55a9
4 changed files with 21 additions and 17 deletions

View File

@ -28,19 +28,22 @@ class AttentionBlock(nn.Module):
return x * attention_weights
class SISUDiscriminator(nn.Module):
def __init__(self, layers=4): #Increased base layer count
def __init__(self, layers=64): #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(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=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 * 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.
)