| 214 | class autoencoder_vgg5(nn.Module): # 36.78 PSNR |
| 215 | ''' vgg encoder with bilinear upsampling''' |
| 216 | def __init__(self): |
| 217 | super(autoencoder_vgg5, self).__init__() |
| 218 | self.encoder = models.vgg19(pretrained=True).features |
| 219 | self.decoder = nn.Sequential( |
| 220 | # (b, 512, 14, 14) |
| 221 | nn.Conv2d(512, 512, 3, stride=1, padding=1), |
| 222 | nn.ReLU(True), |
| 223 | # (b, 512, 28, 28) |
| 224 | nn.Conv2d(512, 512, 3, stride=1, padding=1), |
| 225 | nn.ReLU(True), |
| 226 | # (b, 256, 56, 56) |
| 227 | nn.Conv2d(512, 256, 3, stride=1, padding=1), |
| 228 | nn.ReLU(True), |
| 229 | # (b, 128, 112, 112) |
| 230 | nn.Conv2d(256, 128, 3, stride=1, padding=1), |
| 231 | nn.ReLU(True), |
| 232 | # (b, 64, 224, 224) |
| 233 | nn.Conv2d(128, 64, 3, stride=1, padding=1), |
| 234 | nn.ReLU(True), |
| 235 | nn.Conv2d(64, 3, 3, stride=1, padding=1), |
| 236 | # nn.Tanh() # MSELoss |
| 237 | nn.Sigmoid() # BCELoss |
| 238 | ) |
| 239 | def forward(self, x): |
| 240 | # pdb.set_trace() |
| 241 | feat = [] |