| 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 |