25 lines
934 B
Python
25 lines
934 B
Python
import torch.nn as nn
|
|
|
|
class SISUDiscriminator(nn.Module):
|
|
def __init__(self):
|
|
super(SISUDiscriminator, self).__init__()
|
|
self.model = nn.Sequential(
|
|
nn.Conv1d(2, 128, kernel_size=3, padding=1),
|
|
#nn.LeakyReLU(0.2, inplace=True),
|
|
nn.Conv1d(128, 256, kernel_size=3, padding=1),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
nn.Conv1d(256, 128, kernel_size=3, padding=1),
|
|
#nn.LeakyReLU(0.2, inplace=True),
|
|
nn.Conv1d(128, 64, kernel_size=3, padding=1),
|
|
#nn.LeakyReLU(0.2, inplace=True),
|
|
nn.Conv1d(64, 1, kernel_size=3, padding=1),
|
|
#nn.LeakyReLU(0.2, inplace=True),
|
|
)
|
|
self.global_avg_pool = nn.AdaptiveAvgPool1d(1) # Output size (1,)
|
|
|
|
def forward(self, x):
|
|
x = self.model(x)
|
|
x = self.global_avg_pool(x)
|
|
x = x.view(-1, 1) # Flatten to (batch_size, 1)
|
|
return x
|