| 84 | |
| 85 | |
| 86 | class AuxiliaryHeadImageNet(nn.Module): |
| 87 | |
| 88 | def __init__(self, C, num_classes): |
| 89 | """assuming input size 14x14""" |
| 90 | super(AuxiliaryHeadImageNet, self).__init__() |
| 91 | self.features = nn.Sequential( |
| 92 | nn.ReLU(inplace=True), |
| 93 | nn.AvgPool2d(5, stride=2, padding=0, count_include_pad=False), |
| 94 | nn.Conv2d(C, 128, 1, bias=False), |
| 95 | nn.BatchNorm2d(128), |
| 96 | nn.ReLU(inplace=True), |
| 97 | nn.Conv2d(128, 768, 2, bias=False), |
| 98 | # NOTE: This batchnorm was omitted in my earlier implementation due to a typo. |
| 99 | # Commenting it out for consistency with the experiments in the paper. |
| 100 | # nn.BatchNorm2d(768), |
| 101 | nn.ReLU(inplace=True) |
| 102 | ) |
| 103 | self.classifier = nn.Linear(768, num_classes) |
| 104 | |
| 105 | def forward(self, x): |
| 106 | x = self.features(x) |
| 107 | x = self.classifier(x.view(x.size(0),-1)) |
| 108 | return x |
| 109 | |
| 110 | |
| 111 | class NetworkCIFAR(nn.Module): |