From 909007ee6ab209ceaa1800cea4014ddc9be6bf32 Mon Sep 17 00:00:00 2001 From: James Betker Date: Sun, 5 Jul 2020 21:58:35 -0600 Subject: [PATCH] Add G_warmup Let the Generator get to a point where it is at least competing with the discriminator before firing off. Backwards from most GAN architectures, but this one is a bit different from most. --- codes/models/SRGAN_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/codes/models/SRGAN_model.py b/codes/models/SRGAN_model.py index 7e4fd2e7..fa330a09 100644 --- a/codes/models/SRGAN_model.py +++ b/codes/models/SRGAN_model.py @@ -89,6 +89,7 @@ class SRGANModel(BaseModel): # D_update_ratio and D_init_iters self.D_update_ratio = train_opt['D_update_ratio'] if train_opt['D_update_ratio'] else 1 self.D_init_iters = train_opt['D_init_iters'] if train_opt['D_init_iters'] else 0 + self.G_warmup = train_opt['G_warmup'] if train_opt['G_warmup'] else 0 self.D_noise_theta = train_opt['D_noise_theta_init'] if train_opt['D_noise_theta_init'] else 0 self.D_noise_final = train_opt['D_noise_final_it'] if train_opt['D_noise_final_it'] else 0 self.D_noise_theta_floor = train_opt['D_noise_theta_floor'] if train_opt['D_noise_theta_floor'] else 0 @@ -300,7 +301,7 @@ class SRGANModel(BaseModel): _t = time() # D - if self.l_gan_w > 0: + if self.l_gan_w > 0 and step > self.G_warmup: for p in self.netD.parameters(): p.requires_grad = True @@ -413,7 +414,7 @@ class SRGANModel(BaseModel): if self.l_gan_w > 0: self.add_log_entry('l_g_gan', l_g_gan.item()) self.add_log_entry('l_g_total', l_g_total.item() * self.mega_batch_factor) - if self.l_gan_w > 0: + if self.l_gan_w > 0 and step > self.G_warmup: self.add_log_entry('l_d_real', l_d_real.item() * self.mega_batch_factor) self.add_log_entry('l_d_fake', l_d_fake.item() * self.mega_batch_factor) self.add_log_entry('D_fake', torch.mean(pred_d_fake.detach()))