docstring for FineDecoder
| 251 | |
| 252 | |
| 253 | class FineDecoderV2(nn.Module): |
| 254 | """docstring for FineDecoder""" |
| 255 | |
| 256 | def __init__(self, image_nc, feature_nc, ngf, img_f, layers, num_block, norm_layer=nn.BatchNorm2d, |
| 257 | nonlinearity=nn.LeakyReLU(), use_spect=False): |
| 258 | super(FineDecoderV2, self).__init__() |
| 259 | self.layers = layers |
| 260 | for i in range(layers)[::-1]: |
| 261 | in_channels = min(ngf * (2 ** (i + 1)), img_f) |
| 262 | out_channels = min(ngf * (2 ** i), img_f) |
| 263 | up = UpBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect) |
| 264 | res = FineADAINResBlocks(num_block, in_channels, feature_nc, norm_layer, nonlinearity, use_spect) |
| 265 | jump = Jump(out_channels, norm_layer, nonlinearity, use_spect) |
| 266 | |
| 267 | setattr(self, 'up' + str(i), up) |
| 268 | setattr(self, 'res' + str(i), res) |
| 269 | setattr(self, 'jump' + str(i), jump) |
| 270 | |
| 271 | self.final1 = FinalBlock2d(out_channels, image_nc, use_spect, 'tanh') |
| 272 | self.final2 = FinalBlock2d(out_channels, image_nc, use_spect, 'tanh') |
| 273 | |
| 274 | self.output_nc = out_channels |
| 275 | |
| 276 | def forward(self, x, z): |
| 277 | out = x.pop() |
| 278 | for i in range(self.layers)[::-1]: |
| 279 | res_model = getattr(self, 'res' + str(i)) |
| 280 | up_model = getattr(self, 'up' + str(i)) |
| 281 | jump_model = getattr(self, 'jump' + str(i)) |
| 282 | out = res_model(out, z) |
| 283 | out = up_model(out) |
| 284 | out = jump_model(x.pop()) + out |
| 285 | out_image1 = self.final1(out) |
| 286 | out_image2 = self.final2(out) |
| 287 | return [out_image1, out_image2] |
| 288 | |
| 289 | |
| 290 | class FirstBlock2d(nn.Module): |