Spaces:
Sleeping
Sleeping
| import torch.nn as nn | |
| class CNN(nn.Module): | |
| def __init__(self, num_classes=6): | |
| super(CNN, self).__init__() | |
| self.block1 = self._block(3, 32) | |
| self.block2 = self._block(32, 64) | |
| self.block3 = self._block(64, 128) | |
| self.block4 = self._block(128, 256) | |
| self.gap = nn.AdaptiveAvgPool2d(1) | |
| self.fc1 = nn.Linear(256, 128) | |
| self.fc2 = nn.Linear(128, num_classes) | |
| self.relu = nn.ReLU() | |
| self.dropout = nn.Dropout(0.5) | |
| def _block(self, in_c, out_c): | |
| return nn.Sequential( | |
| nn.Conv2d(in_c, out_c, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(out_c), | |
| nn.ReLU(), | |
| nn.MaxPool2d(2) | |
| ) | |
| def forward(self, x): | |
| x = self.block1(x) | |
| x = self.block2(x) | |
| x = self.block3(x) | |
| x = self.block4(x) | |
| x = self.gap(x) | |
| x = x.view(x.size(0), -1) | |
| x = self.dropout(x) | |
| x = self.relu(self.fc1(x)) | |
| x = self.dropout(x) | |
| x = self.fc2(x) | |
| return x |