Create a Unet-based generator
| 442 | |
| 443 | |
| 444 | class UnetGenerator(nn.Module): |
| 445 | """Create a Unet-based generator""" |
| 446 | |
| 447 | def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False): |
| 448 | """Construct a Unet generator |
| 449 | Parameters: |
| 450 | input_nc (int) -- the number of channels in input images |
| 451 | output_nc (int) -- the number of channels in output images |
| 452 | num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7, |
| 453 | image of size 128x128 will become of size 1x1 # at the bottleneck |
| 454 | ngf (int) -- the number of filters in the last conv layer |
| 455 | norm_layer -- normalization layer |
| 456 | |
| 457 | We construct the U-Net from the innermost layer to the outermost layer. |
| 458 | It is a recursive process. |
| 459 | """ |
| 460 | super(UnetGenerator, self).__init__() |
| 461 | # construct unet structure |
| 462 | unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True) # add the innermost layer |
| 463 | for i in range(num_downs - 5): # add intermediate layers with ngf * 8 filters |
| 464 | unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout) |
| 465 | # gradually reduce the number of filters from ngf * 8 to ngf |
| 466 | unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer) |
| 467 | unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer) |
| 468 | unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer) |
| 469 | self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer) # add the outermost layer |
| 470 | |
| 471 | def forward(self, input): |
| 472 | """Standard forward""" |
| 473 | return self.model(input) |
| 474 | |
| 475 | |
| 476 | class UnetSkipConnectionBlock(nn.Module): |