| 102 | # Deeper neural network class to be used as teacher: |
| 103 | class DeepNN(nn.Module): |
| 104 | def __init__(self, num_classes=10): |
| 105 | super(DeepNN, self).__init__() |
| 106 | self.features = nn.Sequential( |
| 107 | nn.Conv2d(3, 128, kernel_size=3, padding=1), |
| 108 | nn.ReLU(), |
| 109 | nn.Conv2d(128, 64, kernel_size=3, padding=1), |
| 110 | nn.ReLU(), |
| 111 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 112 | nn.Conv2d(64, 64, kernel_size=3, padding=1), |
| 113 | nn.ReLU(), |
| 114 | nn.Conv2d(64, 32, kernel_size=3, padding=1), |
| 115 | nn.ReLU(), |
| 116 | nn.MaxPool2d(kernel_size=2, stride=2), |
| 117 | ) |
| 118 | self.classifier = nn.Sequential( |
| 119 | nn.Linear(2048, 512), |
| 120 | nn.ReLU(), |
| 121 | nn.Dropout(0.1), |
| 122 | nn.Linear(512, num_classes) |
| 123 | ) |
| 124 | |
| 125 | def forward(self, x): |
| 126 | x = self.features(x) |