Class autoencoder_vgg4

script/feature/model.py:157–212

vgg encoder with bilinear upsampling

        return feat_out, x

class autoencoder_vgg4(nn.Module): # 35.54 PSNR 36.05 BCELoss (120x120)
    ''' vgg encoder with bilinear upsampling'''
    def __init__(self):
        super(autoencoder_vgg4, self).__init__()
        self.encoder = models.vgg19(pretrained=True).features
        # receptive field not equal? so maybe it does not work very well
        self.decoder = nn.Sequential(
            # (b, 512, 14, 14)
            # nn.UpsamplingBilinear2d(scale_factor=2), # upsample to feature map's size
            nn.Conv2d(512, 512, 3, stride=1, padding=1),
            nn.ReLU(True),
            # (b, 256, 56, 56)
            # nn.UpsamplingBilinear2d(scale_factor=4),
            nn.Conv2d(512, 256, 3, stride=1, padding=1),
            nn.ReLU(True),
            # (b, 64, 224, 224)
            # nn.UpsamplingBilinear2d(scale_factor=4),
            nn.Conv2d(256, 64, 3, stride=1, padding=1),
            nn.ReLU(True),
            nn.Conv2d(64, 3, 3, stride=1, padding=1),
            # nn.Tanh() # MSELoss
            nn.Sigmoid() # BCELoss
        )
    def forward(self, x):
        # pdb.set_trace()
        feat = []
        feat_out = []
        for i in range(len(self.encoder)):
            # print("layer {} encoder layer: {}".format(i, self.encoder[i]))
            x = self.encoder[i](x)
            if i == 35: # ReLU-36
                feat.append(x)
            elif i == 17: # ReLU-17
                feat.append(x)
            elif i == 3: # ReLU-4
                feat.append(x)      

        for i in range(len(self.decoder)):
            # print("layer {} decoder layer: {}".format(i, self.decoder[i]))
            x = self.decoder[i](x)
            if i == 1:
                _, _, h, w = feat[2].shape
                x = nn.UpsamplingBilinear2d(size=(h,w))(x)
                x = x + feat[2]
                feat_out.append(x)
            elif i == 3:
                _, _, h, w = feat[1].shape
                x = nn.UpsamplingBilinear2d(size=(h,w))(x)
                x = x + feat[1]
                feat_out.append(x)
            elif i == 5:
                _, _, h, w = feat[0].shape
                x = nn.UpsamplingBilinear2d(size=(h,w))(x)
                x = x + feat[0]
                feat_out.append(x)
        return feat_out, x

class autoencoder_vgg5(nn.Module): # 36.78 PSNR