| import torch.nn as nn | |
| class MNISTCNN(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.features = nn.Sequential( | |
| nn.Conv2d(1, 32, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.MaxPool2d(2), | |
| nn.Conv2d(32, 64, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.MaxPool2d(2), | |
| ) | |
| self.classifier = nn.Sequential( | |
| nn.Flatten(), | |
| nn.Linear(64 * 7 * 7, 128), | |
| nn.ReLU(), | |
| nn.Dropout(0.3), | |
| nn.Linear(128, 10), | |
| ) | |
| def forward(self, x): | |
| x = self.features(x) | |
| return self.classifier(x) | |