| 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) | |