| 415 | |
| 416 | # Create a similar student class where we return a tuple. We do not apply pooling after flattening. |
| 417 | class ModifiedLightNNCosine(nn.Module): |
| 418 | def __init__(self, num_classes=10): |
| 419 | super(ModifiedLightNNCosine, self).__init__() |
| 420 | self.features = nn.Sequential( |
| 421 | nn.Conv2d(3, 16, kernel_size=3, padding=1), |
| 422 | nn.ReLU(), |
| 423 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 424 | nn.Conv2d(16, 16, kernel_size=3, padding=1), |
| 425 | nn.ReLU(), |
| 426 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 427 | ) |
| 428 | self.classifier = nn.Sequential( |
| 429 | nn.Linear(1024, 256), |
| 430 | nn.ReLU(), |
| 431 | nn.Dropout(0.1), |
| 432 | nn.Linear(256, num_classes) |
| 433 | ) |
| 434 | |
| 435 | def forward(self, x): |
| 436 | x = self.features(x) |
| 437 | flattened_conv_output = torch.flatten(x, 1) |
| 438 | x = self.classifier(flattened_conv_output) |
| 439 | return x, flattened_conv_output |
| 440 | |
| 441 | # We do not have to train the modified deep network from scratch of course, we just load its weights from the trained instance |
| 442 | modified_nn_deep = ModifiedDeepNNCosine(num_classes=10).to(device) |
no outgoing calls
no test coverage detected