MCPcopy Create free account
hub / github.com/ActiveVisionLab/DFNet / autoencoder_vgg3

Class autoencoder_vgg3

script/feature/model.py:113–155  ·  view source on GitHub ↗

vgg encoder

Source from the content-addressed store, hash-verified

111 return None, x
112
113class 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
157class autoencoder_vgg4(nn.Module): # 35.54 PSNR 36.05 BCELoss (120x120)
158 ''' vgg encoder with bilinear upsampling'''

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected