SISU/discriminator.py
2024-12-25 00:09:57 +02:00

31 lines
1.1 KiB
Python

import torch.nn as nn
class SISUDiscriminator(nn.Module):
def __init__(self):
super(SISUDiscriminator, self).__init__()
layers = 32
self.model = nn.Sequential(
nn.Conv1d(1, layers, kernel_size=5, stride=2, padding=2),
nn.BatchNorm1d(layers),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv1d(layers, layers * 2, kernel_size=5, stride=2, padding=2),
nn.BatchNorm1d(layers * 2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv1d(layers * 2, layers * 4, kernel_size=5, stride=2, padding=2),
nn.BatchNorm1d(layers * 4),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv1d(layers * 4, layers * 8, kernel_size=5, stride=2, padding=2),
nn.BatchNorm1d(layers * 8),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv1d(layers * 8, 1, kernel_size=3, padding=1),
)
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