Defines the Unet submodule with skip connection. X -------------------identity---------------------- |-- downsampling -- |submodule| -- upsampling --|
| 1241 | |
| 1242 | |
| 1243 | class UnetSkipConnectionBlock(nn.Module): |
| 1244 | """Defines the Unet submodule with skip connection. |
| 1245 | X -------------------identity---------------------- |
| 1246 | |-- downsampling -- |submodule| -- upsampling --| |
| 1247 | """ |
| 1248 | |
| 1249 | def __init__(self, outer_nc, inner_nc, input_nc=None, |
| 1250 | submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False): |
| 1251 | """Construct a Unet submodule with skip connections. |
| 1252 | |
| 1253 | Parameters: |
| 1254 | outer_nc (int) -- the number of filters in the outer conv layer |
| 1255 | inner_nc (int) -- the number of filters in the inner conv layer |
| 1256 | input_nc (int) -- the number of channels in input images/features |
| 1257 | submodule (UnetSkipConnectionBlock) -- previously defined submodules |
| 1258 | outermost (bool) -- if this module is the outermost module |
| 1259 | innermost (bool) -- if this module is the innermost module |
| 1260 | norm_layer -- normalization layer |
| 1261 | use_dropout (bool) -- if use dropout layers. |
| 1262 | """ |
| 1263 | super(UnetSkipConnectionBlock, self).__init__() |
| 1264 | self.outermost = outermost |
| 1265 | if type(norm_layer) == functools.partial: |
| 1266 | use_bias = norm_layer.func == nn.InstanceNorm2d |
| 1267 | else: |
| 1268 | use_bias = norm_layer == nn.InstanceNorm2d |
| 1269 | if input_nc is None: |
| 1270 | input_nc = outer_nc |
| 1271 | downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, |
| 1272 | stride=2, padding=1, bias=use_bias) |
| 1273 | downrelu = nn.LeakyReLU(0.2, True) |
| 1274 | downnorm = norm_layer(inner_nc) |
| 1275 | uprelu = nn.ReLU(True) |
| 1276 | upnorm = norm_layer(outer_nc) |
| 1277 | |
| 1278 | if outermost: |
| 1279 | upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, |
| 1280 | kernel_size=4, stride=2, |
| 1281 | padding=1) |
| 1282 | down = [downconv] |
| 1283 | up = [uprelu, upconv, nn.Tanh()] |
| 1284 | model = down + [submodule] + up |
| 1285 | elif innermost: |
| 1286 | upconv = nn.ConvTranspose2d(inner_nc, outer_nc, |
| 1287 | kernel_size=4, stride=2, |
| 1288 | padding=1, bias=use_bias) |
| 1289 | down = [downrelu, downconv] |
| 1290 | up = [uprelu, upconv, upnorm] |
| 1291 | model = down + up |
| 1292 | else: |
| 1293 | upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, |
| 1294 | kernel_size=4, stride=2, |
| 1295 | padding=1, bias=use_bias) |
| 1296 | down = [downrelu, downconv, downnorm] |
| 1297 | up = [uprelu, upconv, upnorm] |
| 1298 | |
| 1299 | if use_dropout: |
| 1300 | model = down + [submodule] + up + [nn.Dropout(0.5)] |