Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class Net(nn.Module): | |
| def __init__(self): | |
| super(Net, self).__init__() | |
| # First Conv Block | |
| self.conv1_1 = nn.Conv2d(3, 32, 3, padding=1) | |
| self.conv1_2 = nn.Conv2d(32, 32, 3, padding=1) | |
| self.pool = nn.MaxPool2d(2, 2) | |
| self.dropout1 = nn.Dropout(0.2) | |
| # Second Conv Block | |
| self.conv2_1 = nn.Conv2d(32, 64, 3, padding=1) | |
| self.conv2_2 = nn.Conv2d(64, 64, 3, padding=1) | |
| self.dropout2 = nn.Dropout(0.3) | |
| # Third Conv Block | |
| self.conv3_1 = nn.Conv2d(64, 128, 3, padding=1) | |
| self.conv3_2 = nn.Conv2d(128, 128, 3, padding=1) | |
| self.dropout3 = nn.Dropout(0.4) | |
| # Dense Layers | |
| self.fc1 = nn.Linear(128 * 4 * 4, 128) | |
| self.dropout4 = nn.Dropout(0.5) | |
| self.fc2 = nn.Linear(128, 10) | |
| def forward(self, x): | |
| # Block 1 | |
| x = F.relu(self.conv1_1(x)) | |
| x = F.relu(self.conv1_2(x)) | |
| x = self.pool(x) | |
| x = self.dropout1(x) | |
| # Block 2 | |
| x = F.relu(self.conv2_1(x)) | |
| x = F.relu(self.conv2_2(x)) | |
| x = self.pool(x) | |
| x = self.dropout2(x) | |
| # Block 3 | |
| x = F.relu(self.conv3_1(x)) | |
| x = F.relu(self.conv3_2(x)) | |
| x = self.pool(x) | |
| x = self.dropout3(x) | |
| # Flatten | |
| x = x.view(-1, 128 * 4 * 4) | |
| # Dense | |
| x = F.relu(self.fc1(x)) | |
| x = self.dropout4(x) | |
| x = self.fc2(x) | |
| return x | |
| def create_model(): | |
| return Net() | |