From 35ab7ecd510f1a9532baabae9bc7e6ae37e3e94c Mon Sep 17 00:00:00 2001 From: Ray Wong Date: Sat, 22 Aug 2020 20:57:03 +0800 Subject: [PATCH] update cycle module --- model/GAN/residual_generator.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/model/GAN/residual_generator.py b/model/GAN/residual_generator.py index 9c4adac..ea4abbc 100644 --- a/model/GAN/residual_generator.py +++ b/model/GAN/residual_generator.py @@ -118,17 +118,16 @@ class ResGenerator(nn.Module): multiple = 2 ** i submodules += [ nn.Conv2d(in_channels=base_channels * multiple, out_channels=base_channels * multiple * 2, - kernel_size=3, - stride=2, padding=1, bias=use_bias), + kernel_size=3, stride=2, padding=1, bias=use_bias), norm_layer(num_features=base_channels * multiple * 2), nn.ReLU(inplace=True) ] self.encoder = nn.Sequential(*submodules) res_block_channels = num_down_sampling ** 2 * base_channels - self.res_blocks = nn.ModuleList( - [ResidualBlock(res_block_channels, padding_mode, norm_type, use_dropout=use_dropout) for _ in - range(num_blocks)]) + self.resnet_middle = nn.Sequential( + *[ResidualBlock(res_block_channels, padding_mode, norm_type, use_dropout=use_dropout) for _ in + range(num_blocks)]) # up sampling submodules = [] @@ -149,14 +148,13 @@ class ResGenerator(nn.Module): def forward(self, x): x = self.encoder(self.start_conv(x)) - for rb in self.res_blocks: - x = rb(x) + x = self.resnet_middle(x) return self.end_conv(self.decoder(x)) @MODEL.register_module() class PatchDiscriminator(nn.Module): - def __init__(self, in_channels, base_channels=64, num_conv=3, norm_type="BN"): + def __init__(self, in_channels, base_channels=64, num_conv=3, norm_type="IN"): super(PatchDiscriminator, self).__init__() assert num_conv >= 0, f'Number of conv blocks must be non-negative, but got {num_conv}.' norm_layer = _select_norm_layer(norm_type)