import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1), ) self.fc = nn.Linear(128, 2) # Negative=0, Positive=1 def forward(self, x): x = self.net(x) x = x.view(x.size(0), -1) return self.fc(x)