| 120 | return x |
| 121 | |
| 122 | class ResNet(nn.Module): |
| 123 | def __init__(self, block, num_blocks, n_class=10, bayesian=False): |
| 124 | super(ResNet, self).__init__() |
| 125 | # self.in_planes = 16 |
| 126 | self.embDim = 128 * block.expansion |
| 127 | # self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False) |
| 128 | # self.bn1 = nn.BatchNorm2d(16) |
| 129 | # self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1) |
| 130 | # self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2) |
| 131 | # self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2) |
| 132 | # self.layer4 = self._make_layer(block, 128, num_blocks[3], stride=2) |
| 133 | # self.linear = nn.Linear(128 * block.expansion, n_class) |
| 134 | |
| 135 | # self.dis_fc1 = nn.Linear(512, 50) |
| 136 | # self.dis_fc2 = nn.Linear(50, 1) |
| 137 | |
| 138 | self.feature_extractor = resnet_fea(block, num_blocks) |
| 139 | self.linear = resnet_clf(block, n_class) |
| 140 | self.discriminator = resnet_dis(self.embDim) |
| 141 | self.bayesian = bayesian |
| 142 | |
| 143 | # def _make_layer(self, block, planes, num_blocks, stride): |
| 144 | # strides = [stride] + [1]*(num_blocks-1) |
| 145 | # layers = [] |
| 146 | # for stride in strides: |
| 147 | # layers.append(block(self.in_planes, planes, stride)) |
| 148 | # self.in_planes = planes * block.expansion |
| 149 | # return nn.Sequential(*layers) |
| 150 | |
| 151 | # def feature_extractor(self, x): # feature extractor |
| 152 | # out = F.relu(self.bn1(self.conv1(x))) |
| 153 | # out = self.layer1(out) |
| 154 | # out = self.layer2(out) |
| 155 | # out = self.layer3(out) |
| 156 | # out = self.layer4(out) |
| 157 | # out = F.avg_pool2d(out, 4) |
| 158 | # emb = out.view(out.size(0), -1) |
| 159 | # return emb |
| 160 | |
| 161 | |
| 162 | def forward(self, x, intermediate=False): |
| 163 | out, in_values = self.feature_extractor(x, x.shape[2]) |
| 164 | # apply dropout to approximate the bayesian networks |
| 165 | out = F.dropout(out, p=0.2, training=self.bayesian) |
| 166 | # emb = emb.view(emb.size(0), -1) |
| 167 | out, emb = self.linear(out) |
| 168 | if intermediate == True: |
| 169 | return out, emb, in_values |
| 170 | else: |
| 171 | return out, emb |
| 172 | |
| 173 | def get_embedding_dim(self): |
| 174 | return self.embDim |
| 175 | |
| 176 | |
| 177 | def ResNet18(n_class, bayesian=False): |