File size: 559 Bytes
56b32ab a667a5b 56b32ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
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)
|