Downsampling block for use in encoder.
| 245 | return out_image |
| 246 | |
| 247 | class FirstBlock2d(nn.Module): |
| 248 | """ |
| 249 | Downsampling block for use in encoder. |
| 250 | """ |
| 251 | def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False): |
| 252 | super(FirstBlock2d, self).__init__() |
| 253 | kwargs = {'kernel_size': 7, 'stride': 1, 'padding': 3} |
| 254 | conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect) |
| 255 | |
| 256 | if type(norm_layer) == type(None): |
| 257 | self.model = nn.Sequential(conv, nonlinearity) |
| 258 | else: |
| 259 | self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity) |
| 260 | |
| 261 | |
| 262 | def forward(self, x): |
| 263 | out = self.model(x) |
| 264 | return out |
| 265 | |
| 266 | class DownBlock2d(nn.Module): |
| 267 | def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False): |