vgg encoder with bilinear upsampling
| 293 | return feat_out, x |
| 294 | |
| 295 | class autoencoder_vgg6(nn.Module): # robust feature extractors |
| 296 | ''' vgg encoder with bilinear upsampling''' |
| 297 | def __init__(self): |
| 298 | super(autoencoder_vgg6, self).__init__() |
| 299 | self.encoder = models.vgg19(pretrained=True).features |
| 300 | self.decoder = nn.Sequential( |
| 301 | # (b, 512, 14, 14) |
| 302 | nn.Conv2d(512, 512, 3, stride=1, padding=1), |
| 303 | nn.ReLU(True), |
| 304 | # (b, 512, 28, 28) |
| 305 | nn.Conv2d(512, 512, 3, stride=1, padding=1), |
| 306 | nn.ReLU(True), |
| 307 | # (b, 256, 56, 56) |
| 308 | nn.Conv2d(512, 256, 3, stride=1, padding=1), |
| 309 | nn.ReLU(True), |
| 310 | # (b, 128, 112, 112) |
| 311 | nn.Conv2d(256, 128, 3, stride=1, padding=1), |
| 312 | nn.ReLU(True), |
| 313 | # (b, 64, 224, 224) |
| 314 | nn.Conv2d(128, 64, 3, stride=1, padding=1), |
| 315 | nn.ReLU(True), |
| 316 | # nn.Conv2d(64, 3, 3, stride=1, padding=1), |
| 317 | # nn.Tanh() # MSELoss |
| 318 | # nn.Sigmoid() # BCELoss |
| 319 | ) |
| 320 | |
| 321 | def forward(self, x, upsampleH, upsampleW): # |
| 322 | feat = [] |
| 323 | feat_out = [] # we only use high level features |
| 324 | for i in range(len(self.encoder)): |
| 325 | # print("layer {} encoder layer: {}".format(i, self.encoder[i])) |
| 326 | x = self.encoder[i](x) |
| 327 | if i == 3: # ReLU-4 |
| 328 | feat.append(x) |
| 329 | elif i == 8: # ReLU-9 |
| 330 | feat.append(x) |
| 331 | elif i == 17: # ReLU-18 |
| 332 | feat.append(x) |
| 333 | elif i == 26: # ReLU-27 |
| 334 | feat.append(x) |
| 335 | elif i == 35: # ReLU-36 |
| 336 | feat.append(x) |
| 337 | |
| 338 | for i in range(len(self.decoder)): |
| 339 | # print("layer {} decoder layer: {}".format(i, self.decoder[i])) |
| 340 | x = self.decoder[i](x) |
| 341 | if i == 1: |
| 342 | _, _, h, w = feat[4].shape |
| 343 | x = nn.UpsamplingBilinear2d(size=(h,w))(x) |
| 344 | x = x + feat[4] |
| 345 | elif i == 3: |
| 346 | _, _, h, w = feat[3].shape |
| 347 | x = nn.UpsamplingBilinear2d(size=(h,w))(x) |
| 348 | x = x + feat[3] |
| 349 | elif i == 5: |
| 350 | _, _, h, w = feat[2].shape |
| 351 | x = nn.UpsamplingBilinear2d(size=(h,w))(x) |
| 352 | x = x + feat[2] |
nothing calls this directly
no outgoing calls
no test coverage detected