From 0a9b85f2399369ed6dfafc3ecd57c26ea7493dcb Mon Sep 17 00:00:00 2001 From: James Betker Date: Mon, 31 Aug 2020 09:50:30 -0600 Subject: [PATCH] Fix vgg_gn input_img_factor --- codes/models/archs/discriminator_vgg_arch.py | 2 +- codes/models/networks.py | 4 ++-- codes/train.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/codes/models/archs/discriminator_vgg_arch.py b/codes/models/archs/discriminator_vgg_arch.py index 6ba39264..d9f7f29f 100644 --- a/codes/models/archs/discriminator_vgg_arch.py +++ b/codes/models/archs/discriminator_vgg_arch.py @@ -108,7 +108,7 @@ class Discriminator_VGG_128_GN(nn.Module): self.bn4_1 = nn.GroupNorm(8, nf * 8, affine=True) final_nf = nf * 8 - self.linear1 = nn.Linear(final_nf * 4 * input_img_factor * 4 * input_img_factor, 100) + self.linear1 = nn.Linear(int(final_nf * 4 * input_img_factor * 4 * input_img_factor), 100) self.linear2 = nn.Linear(100, 1) # activation function diff --git a/codes/models/networks.py b/codes/models/networks.py index 31dbce05..fd16e950 100644 --- a/codes/models/networks.py +++ b/codes/models/networks.py @@ -160,9 +160,9 @@ def define_D_net(opt_net, img_sz=None, wrap=False): which_model = opt_net['which_model_D'] if which_model == 'discriminator_vgg_128': - netD = SRGAN_arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], nf=opt_net['nf'], input_img_factor=img_sz // 128, extra_conv=opt_net['extra_conv']) + netD = SRGAN_arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], nf=opt_net['nf'], input_img_factor=img_sz / 128, extra_conv=opt_net['extra_conv']) elif which_model == 'discriminator_vgg_128_gn': - netD = SRGAN_arch.Discriminator_VGG_128_GN(in_nc=opt_net['in_nc'], nf=opt_net['nf'], input_img_factor=img_sz // 128) + netD = SRGAN_arch.Discriminator_VGG_128_GN(in_nc=opt_net['in_nc'], nf=opt_net['nf'], input_img_factor=img_sz / 128) if wrap: netD = GradDiscWrapper(netD) elif which_model == 'discriminator_resnet': diff --git a/codes/train.py b/codes/train.py index 82ae665f..661e9b92 100644 --- a/codes/train.py +++ b/codes/train.py @@ -32,7 +32,7 @@ def init_dist(backend='nccl', **kwargs): def main(): #### options parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/srgan_compute_feature.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/pretrain_spsr_switched2_psnr.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0)