| 147 | |
| 148 | # LeNet Model definition |
| 149 | class Net(nn.Module): |
| 150 | def __init__(self): |
| 151 | super(Net, self).__init__() |
| 152 | self.conv1 = nn.Conv2d(1, 32, 3, 1) |
| 153 | self.conv2 = nn.Conv2d(32, 64, 3, 1) |
| 154 | self.dropout1 = nn.Dropout(0.25) |
| 155 | self.dropout2 = nn.Dropout(0.5) |
| 156 | self.fc1 = nn.Linear(9216, 128) |
| 157 | self.fc2 = nn.Linear(128, 10) |
| 158 | |
| 159 | def forward(self, x): |
| 160 | x = self.conv1(x) |
| 161 | x = F.relu(x) |
| 162 | x = self.conv2(x) |
| 163 | x = F.relu(x) |
| 164 | x = F.max_pool2d(x, 2) |
| 165 | x = self.dropout1(x) |
| 166 | x = torch.flatten(x, 1) |
| 167 | x = self.fc1(x) |
| 168 | x = F.relu(x) |
| 169 | x = self.dropout2(x) |
| 170 | x = self.fc2(x) |
| 171 | output = F.log_softmax(x, dim=1) |
| 172 | return output |
| 173 | |
| 174 | # MNIST Test dataset and dataloader declaration |
| 175 | test_loader = torch.utils.data.DataLoader( |