| 385 | # Now, the forward function returns not only the logits of the network but also the flattened hidden representation after the convolutional layer. We include the aforementioned pooling for the modified teacher. |
| 386 | |
| 387 | class ModifiedDeepNNCosine(nn.Module): |
| 388 | def __init__(self, num_classes=10): |
| 389 | super(ModifiedDeepNNCosine, self).__init__() |
| 390 | self.features = nn.Sequential( |
| 391 | nn.Conv2d(3, 128, kernel_size=3, padding=1), |
| 392 | nn.ReLU(), |
| 393 | nn.Conv2d(128, 64, kernel_size=3, padding=1), |
| 394 | nn.ReLU(), |
| 395 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 396 | nn.Conv2d(64, 64, kernel_size=3, padding=1), |
| 397 | nn.ReLU(), |
| 398 | nn.Conv2d(64, 32, kernel_size=3, padding=1), |
| 399 | nn.ReLU(), |
| 400 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 401 | ) |
| 402 | self.classifier = nn.Sequential( |
| 403 | nn.Linear(2048, 512), |
| 404 | nn.ReLU(), |
| 405 | nn.Dropout(0.1), |
| 406 | nn.Linear(512, num_classes) |
| 407 | ) |
| 408 | |
| 409 | def forward(self, x): |
| 410 | x = self.features(x) |
| 411 | flattened_conv_output = torch.flatten(x, 1) |
| 412 | x = self.classifier(flattened_conv_output) |
| 413 | flattened_conv_output_after_pooling = torch.nn.functional.avg_pool1d(flattened_conv_output, 2) |
| 414 | return x, flattened_conv_output_after_pooling |
| 415 | |
| 416 | # Create a similar student class where we return a tuple. We do not apply pooling after flattening. |
| 417 | class ModifiedLightNNCosine(nn.Module): |
no outgoing calls
no test coverage detected