| 25 | 'TinyImageNet': [3, 200], |
| 26 | } |
| 27 | class CNN(nn.Module): |
| 28 | def __init__(self, CNN_name, dataset, dropout=False): |
| 29 | super(CNN, self).__init__() |
| 30 | self.dataset = dataset |
| 31 | self.query_num = 0 |
| 32 | self.features = self._make_layers(mycfg[CNN_name]) |
| 33 | if dropout: |
| 34 | self.classifier = nn.Sequential( |
| 35 | nn.Dropout(0.6), |
| 36 | nn.Linear(512, 256), |
| 37 | nn.ReLU(True), |
| 38 | nn.Linear(256, parameters[self.dataset][1]) ) |
| 39 | else: |
| 40 | self.classifier = nn.Sequential( |
| 41 | nn.Linear(512, 256), |
| 42 | nn.ReLU(True), |
| 43 | nn.Linear(256, parameters[self.dataset][1]) ) |
| 44 | |
| 45 | def forward(self, x): |
| 46 | self.query_num += 1 |
| 47 | out = self.features(x) |
| 48 | out = out.view(out.size(0), -1) |
| 49 | out = self.classifier(out) |
| 50 | return out |
| 51 | |
| 52 | def _make_layers(self, cfg): |
| 53 | layers = [] |
| 54 | in_channels = 3 |
| 55 | for x in cfg: |
| 56 | if x == 'M': |
| 57 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] |
| 58 | else: |
| 59 | layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1), |
| 60 | nn.BatchNorm2d(x, track_running_stats=True), |
| 61 | nn.ReLU(inplace=True)] |
| 62 | in_channels = x |
| 63 | layers += [nn.AvgPool2d(kernel_size=parameters[self.dataset][0], stride=parameters[self.dataset][0])] |
| 64 | return nn.Sequential(*layers) |
| 65 | |
| 66 | class MemGuard(nn.Module): |
| 67 | def __init__(self): |