| 41 | |
| 42 | |
| 43 | class ADAINEncoder(nn.Module): |
| 44 | def __init__(self, image_nc, pose_nc, ngf, img_f, layers, nonlinearity=nn.LeakyReLU(), use_spect=False): |
| 45 | super(ADAINEncoder, self).__init__() |
| 46 | self.layers = layers |
| 47 | self.input_layer = nn.Conv2d(image_nc, ngf, kernel_size=7, stride=1, padding=3) |
| 48 | for i in range(layers): |
| 49 | in_channels = min(ngf * (2**i), img_f) |
| 50 | out_channels = min(ngf *(2**(i+1)), img_f) |
| 51 | model = ADAINEncoderBlock(in_channels, out_channels, pose_nc, nonlinearity, use_spect) |
| 52 | setattr(self, 'encoder' + str(i), model) |
| 53 | self.output_nc = out_channels |
| 54 | |
| 55 | def forward(self, x, z): |
| 56 | out = self.input_layer(x) |
| 57 | out_list = [out] |
| 58 | for i in range(self.layers): |
| 59 | model = getattr(self, 'encoder' + str(i)) |
| 60 | out = model(out, z) |
| 61 | out_list.append(out) |
| 62 | return out_list |
| 63 | |
| 64 | class ADAINDecoder(nn.Module): |
| 65 | """docstring for ADAINDecoder""" |