vgg encoder
| 111 | return None, x |
| 112 | |
| 113 | class autoencoder_vgg3(nn.Module): # psnr: 37.77 |
| 114 | ''' vgg encoder ''' |
| 115 | def __init__(self): |
| 116 | super(autoencoder_vgg3, self).__init__() |
| 117 | self.encoder = models.vgg19(pretrained=True).features |
| 118 | # receptive field not equal? so maybe it does not work very well |
| 119 | self.decoder = nn.Sequential( |
| 120 | nn.ConvTranspose2d(512, 512, 2, stride=2), # (b, 512, 14, 14) |
| 121 | nn.ReLU(True), |
| 122 | nn.ConvTranspose2d(512, 256, 4, stride=4), # (b, 256, 56, 56) |
| 123 | nn.ReLU(True), |
| 124 | nn.ConvTranspose2d(256, 64, 4, stride=4), # (b, 64, 224, 224) |
| 125 | nn.ReLU(True), |
| 126 | nn.Conv2d(64, 3, 3, stride=1, padding=1), |
| 127 | nn.Tanh() # MSELoss |
| 128 | # nn.Sigmoid() # BCELoss |
| 129 | ) |
| 130 | |
| 131 | def forward(self, x): |
| 132 | feat = [] |
| 133 | feat_out = [] |
| 134 | for i in range(len(self.encoder)): |
| 135 | # print("layer {} encoder layer: {}".format(i, self.encoder[i])) |
| 136 | x = self.encoder[i](x) |
| 137 | if i == 35: # ReLU-36 |
| 138 | feat.append(x) |
| 139 | elif i == 17: # ReLU-17 |
| 140 | feat.append(x) |
| 141 | elif i == 3: # ReLU-4 |
| 142 | feat.append(x) |
| 143 | for i in range(len(self.decoder)): |
| 144 | # print("layer {} decoder layer: {}".format(i, self.decoder[i])) |
| 145 | x = self.decoder[i](x) |
| 146 | if i == 1: |
| 147 | x = x + feat[2] |
| 148 | feat_out.append(x) |
| 149 | elif i == 3: |
| 150 | x = x + feat[1] |
| 151 | feat_out.append(x) |
| 152 | elif i == 5: |
| 153 | x = x + feat[0] |
| 154 | feat_out.append(x) |
| 155 | return feat_out, x |
| 156 | |
| 157 | class autoencoder_vgg4(nn.Module): # 35.54 PSNR 36.05 BCELoss (120x120) |
| 158 | ''' vgg encoder with bilinear upsampling''' |
nothing calls this directly
no outgoing calls
no test coverage detected