Merge pull request #24 from DAA233/fix-gradient-penalty-bug

fix gradient penalty bug for PyTorch 1.2+ and update the trainer code

Author: daa233

trainer.py

@@ -4,12 +4,13 @@
 from torch import autograd
 from model.networks import Generator, LocalDis, GlobalDis
 
+
 from utils.tools import get_model_list, local_patch, spatial_discounting_mask
 from utils.logger import get_logger
 
-
 logger = get_logger()
 
+
 class Trainer(nn.Module):
     def __init__(self, config):
         super(Trainer, self).__init__()
@@ -33,6 +34,7 @@ def __init__(self, config):
 
     def forward(self, x, bboxes, masks, ground_truth, compute_loss_g=False):
         self.train()
+        l1_loss = nn.L1Loss()
         losses = {}
 
         x1, x2, offset_flow = self.netG(x, masks)
@@ -42,35 +44,37 @@ def forward(self, x, bboxes, masks, ground_truth, compute_loss_g=False):
         local_patch_x1_inpaint = local_patch(x1_inpaint, bboxes)
         local_patch_x2_inpaint = local_patch(x2_inpaint, bboxes)
 
-        ## D part
+        # D part
         # wgan d loss
-        local_patch_real_pred, local_patch_fake_pred = \
-            self.dis_forward(self.localD, local_patch_gt, local_patch_x2_inpaint.detach())
-        global_real_pred, global_fake_pred = \
-            self.dis_forward(self.globalD, ground_truth, x2_inpaint.detach())
-        losses['wgan_d'] = torch.mean(local_patch_fake_pred - local_patch_real_pred) \
-                           + torch.mean(global_fake_pred - global_real_pred) * self.config['global_wgan_loss_alpha']
+        local_patch_real_pred, local_patch_fake_pred = self.dis_forward(
+            self.localD, local_patch_gt, local_patch_x2_inpaint.detach())
+        global_real_pred, global_fake_pred = self.dis_forward(
+            self.globalD, ground_truth, x2_inpaint.detach())
+        losses['wgan_d'] = torch.mean(local_patch_fake_pred - local_patch_real_pred) + \
+            torch.mean(global_fake_pred - global_real_pred) * self.config['global_wgan_loss_alpha']
         # gradients penalty loss
-        local_penalty = self.calc_gradient_penalty(self.localD, local_patch_gt, local_patch_x2_inpaint.detach())
+        local_penalty = self.calc_gradient_penalty(
+            self.localD, local_patch_gt, local_patch_x2_inpaint.detach())
         global_penalty = self.calc_gradient_penalty(self.globalD, ground_truth, x2_inpaint.detach())
         losses['wgan_gp'] = local_penalty + global_penalty
 
-        ## G part
+        # G part
         if compute_loss_g:
             sd_mask = spatial_discounting_mask(self.config)
-            losses['l1'] = nn.L1Loss()(local_patch_x1_inpaint * sd_mask, local_patch_gt * sd_mask) \
-                           * self.config['coarse_l1_alpha'] \
-                           + nn.L1Loss()(local_patch_x2_inpaint * sd_mask, local_patch_gt * sd_mask)
-            losses['ae'] = nn.L1Loss()(x1 * (1. - masks), ground_truth * (1. - masks)) \
-                           * self.config['coarse_l1_alpha'] \
-                           + nn.L1Loss()(x2 * (1. - masks), ground_truth * (1. - masks))
-            # wgan g loss
-            local_patch_real_pred, local_patch_fake_pred = \
-                self.dis_forward(self.localD, local_patch_gt, local_patch_x2_inpaint)
-            global_real_pred, global_fake_pred = self.dis_forward(self.globalD, ground_truth, x2_inpaint)
+            losses['l1'] = l1_loss(local_patch_x1_inpaint * sd_mask, local_patch_gt * sd_mask) * \
+                self.config['coarse_l1_alpha'] + \
+                l1_loss(local_patch_x2_inpaint * sd_mask, local_patch_gt * sd_mask)
+            losses['ae'] = l1_loss(x1 * (1. - masks), ground_truth * (1. - masks)) * \
+                self.config['coarse_l1_alpha'] + \
+                l1_loss(x2 * (1. - masks), ground_truth * (1. - masks))
 
-            losses['wgan_g'] = - torch.mean(local_patch_fake_pred) \
-                               - torch.mean(global_fake_pred) * self.config['global_wgan_loss_alpha']
+            # wgan g loss
+            local_patch_real_pred, local_patch_fake_pred = self.dis_forward(
+                self.localD, local_patch_gt, local_patch_x2_inpaint)
+            global_real_pred, global_fake_pred = self.dis_forward(
+                self.globalD, ground_truth, x2_inpaint)
+            losses['wgan_g'] = - torch.mean(local_patch_fake_pred) - \
+                torch.mean(global_fake_pred) * self.config['global_wgan_loss_alpha']
 
         return losses, x2_inpaint, offset_flow
 
@@ -85,26 +89,26 @@ def dis_forward(self, netD, ground_truth, x_inpaint):
 
     # Calculate gradient penalty
     def calc_gradient_penalty(self, netD, real_data, fake_data):
-        batch_size, channel, height, width = real_data.size()
-        alpha = torch.rand(batch_size, 1)
-        alpha = alpha.expand(batch_size, int(real_data.nelement() // batch_size)).contiguous() \
-            .view(batch_size, channel, height, width)
+        batch_size = real_data.size(0)
+        alpha = torch.rand(batch_size, 1, 1, 1)
+        alpha = alpha.expand_as(real_data)
         if self.use_cuda:
             alpha = alpha.cuda()
 
-        interpolates = alpha * real_data + ((1 - alpha) * fake_data)
-        interpolates = autograd.Variable(interpolates, requires_grad=True)
+        interpolates = alpha * real_data + (1 - alpha) * fake_data
+        interpolates = interpolates.requires_grad_().clone()
 
         disc_interpolates = netD(interpolates)
-
         grad_outputs = torch.ones(disc_interpolates.size())
+
         if self.use_cuda:
             grad_outputs = grad_outputs.cuda()
+
         gradients = autograd.grad(outputs=disc_interpolates, inputs=interpolates,
-                                  grad_outputs=grad_outputs,
-                                  create_graph=True, retain_graph=True, only_inputs=True)[0]
-        gradients = gradients.view(gradients.size(0), -1)
+                                  grad_outputs=grad_outputs, create_graph=True,
+                                  retain_graph=True, only_inputs=True)[0]
 
+        gradients = gradients.view(batch_size, -1)
         gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
 
         return gradient_penalty
@@ -123,8 +127,10 @@ def save_model(self, checkpoint_dir, iteration):
         dis_name = os.path.join(checkpoint_dir, 'dis_%08d.pt' % iteration)
         opt_name = os.path.join(checkpoint_dir, 'optimizer.pt')
         torch.save(self.netG.state_dict(), gen_name)
-        torch.save({'localD': self.localD.state_dict(), 'globalD': self.globalD.state_dict()}, dis_name)
-        torch.save({'gen': self.optimizer_g.state_dict(), 'dis': self.optimizer_d.state_dict()}, opt_name)
+        torch.save({'localD': self.localD.state_dict(),
+                    'globalD': self.globalD.state_dict()}, dis_name)
+        torch.save({'gen': self.optimizer_g.state_dict(),
+                    'dis': self.optimizer_d.state_dict()}, opt_name)
 
     def resume(self, checkpoint_dir, iteration=0, test=False):
         # Load generators