| 131 | # Lightweight neural network class to be used as student: |
| 132 | class LightNN(nn.Module): |
| 133 | def __init__(self, num_classes=10): |
| 134 | super(LightNN, self).__init__() |
| 135 | self.features = nn.Sequential( |
| 136 | nn.Conv2d(3, 16, kernel_size=3, padding=1), |
| 137 | nn.ReLU(), |
| 138 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 139 | nn.Conv2d(16, 16, kernel_size=3, padding=1), |
| 140 | nn.ReLU(), |
| 141 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 142 | ) |
| 143 | self.classifier = nn.Sequential( |
| 144 | nn.Linear(1024, 256), |
| 145 | nn.ReLU(), |
| 146 | nn.Dropout(0.1), |
| 147 | nn.Linear(256, num_classes) |
| 148 | ) |
| 149 | |
| 150 | def forward(self, x): |
| 151 | x = self.features(x) |