Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations. We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
| 321 | |
| 322 | |
| 323 | class ResnetGenerator(nn.Module): |
| 324 | """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations. |
| 325 | |
| 326 | We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style) |
| 327 | """ |
| 328 | |
| 329 | def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'): |
| 330 | """Construct a Resnet-based generator |
| 331 | |
| 332 | Parameters: |
| 333 | input_nc (int) -- the number of channels in input images |
| 334 | output_nc (int) -- the number of channels in output images |
| 335 | ngf (int) -- the number of filters in the last conv layer |
| 336 | norm_layer -- normalization layer |
| 337 | use_dropout (bool) -- if use dropout layers |
| 338 | n_blocks (int) -- the number of ResNet blocks |
| 339 | padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero |
| 340 | """ |
| 341 | assert(n_blocks >= 0) |
| 342 | super(ResnetGenerator, self).__init__() |
| 343 | if type(norm_layer) == functools.partial: |
| 344 | use_bias = norm_layer.func == nn.InstanceNorm2d |
| 345 | else: |
| 346 | use_bias = norm_layer == nn.InstanceNorm2d |
| 347 | |
| 348 | model = [nn.ReflectionPad2d(3), |
| 349 | nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), |
| 350 | norm_layer(ngf), |
| 351 | nn.ReLU(True)] |
| 352 | |
| 353 | n_downsampling = 2 |
| 354 | for i in range(n_downsampling): # add downsampling layers |
| 355 | mult = 2 ** i |
| 356 | model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias), |
| 357 | norm_layer(ngf * mult * 2), |
| 358 | nn.ReLU(True)] |
| 359 | |
| 360 | mult = 2 ** n_downsampling |
| 361 | for i in range(n_blocks): # add ResNet blocks |
| 362 | |
| 363 | model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)] |
| 364 | |
| 365 | for i in range(n_downsampling): # add upsampling layers |
| 366 | mult = 2 ** (n_downsampling - i) |
| 367 | model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), |
| 368 | kernel_size=3, stride=2, |
| 369 | padding=1, output_padding=1, |
| 370 | bias=use_bias), |
| 371 | norm_layer(int(ngf * mult / 2)), |
| 372 | nn.ReLU(True)] |
| 373 | model += [nn.ReflectionPad2d(3)] |
| 374 | model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)] |
| 375 | model += [nn.Tanh()] |
| 376 | |
| 377 | self.model = nn.Sequential(*model) |
| 378 | |
| 379 | def forward(self, input): |
| 380 | """Standard forward""" |