Training broken after Lightning upgrade

[simonjaq]

Hello all

I developed a quite complex Image2Image translation GAN.
Haven't worked on it for a while. Now I wanted to pick up training it and make some slight adjustments to the network.
In the meantime I had to upgrade Lightning to most recent version because of another project.
I can't get the training to pick up from where it was. Loading a checkpoint and doing inference still works.

The original code was developed under 0.4 version (I think?)

I guess I figured out the changes neccessary for hyperparameter handling. Training does start but loss is NaN right from the beginning.
Here's my current code (shortened). Would be great if one of the pros could guide me at what has to be changed to make it work again.

(Sorry, code is very messy)

## Train Loop

# Train Discriminator and Generator
class Im2Im(pl.LightningModule):
    def __init__(self, *args, **kwargs):
        super(Im2Im, self).__init__()
        
        self.save_hyperparameters()
        #self.hparams = hparams        

        # here I initialize all the networks ... cut
        self.netG_A2B = Unet(backbone_name='vgg19',

 ... snip
        
        self.netG_B2A = Unet(backbone_name='vgg19',


... snip
              
        self.netD_A = MultiscaleDiscriminator(

... snip
        
        self.netD_B = MultiscaleDiscriminator(

... snip
                    
        self.criterion = CharbonnierLoss()

... all the losses initialized ... snip
        
    def forward(self, x):
        
        x, features = self.netG_A2B(x)
        
        return x, features
    
    def training_step(self, batch, batch_nb, optimizer_idx):  # 
        
        if optimizer_idx == 0:
        
            ###### Generators A2B and B2A ######

            #Load data
            self.real_A = batch['A']
            self.real_B = batch['B']
            self.real_A_prev = batch['A_prev']
            self.real_B_prev = batch['B_prev']

            # GAN loss
            self.fake_B, features = self.netG_A2B(self.real_A)  
            #self.fake_AB = torch.cat([real_A,self.fake_B], dim=1)
            self.fake_AB = torch.cat([self.fake_B,self.fake_B], dim=1) # unconditional
            pred_fake = self.netD_B(self.fake_AB)
            loss_GAN_A2B = self.criterion_GAN(pred_fake, True)
        
            self.fake_A, _ = self.netG_B2A(self.real_B)
            #self.fake_BA = torch.cat([real_B,self.fake_A], dim=1)
            self.fake_BA = torch.cat([self.fake_A,self.fake_A], dim=1)
            pred_fake = self.netD_A(self.fake_BA)
            loss_GAN_B2A = self.criterion_GAN(pred_fake, True) 

            # Cycle loss
            recovered_A, _ = self.netG_B2A(self.fake_B)
            loss_cycle_ABA = self.criterion_cycle(recovered_A, self.real_A)*50.0

            recovered_B, _ = self.netG_A2B(self.fake_A)
            loss_cycle_BAB = self.criterion_cycle(recovered_B, self.real_B)*50.0

            #Temporal loss
            flo_bwd = FlowNet(self.real_A, self.real_A_prev)
            flo_fwd = FlowNet(self.real_A_prev, self.real_A)
            occ = occlusion_mask_from_flow(flo_fwd, flo_bwd)
            fake_B_prev, features_prev  = self.netG_A2B(self.real_A_prev)

            t_loss = output_temporal_loss(self.real_A,
                                          self.real_A_prev,
                                          self.fake_B,
                                          fake_B_prev.detach(),
                                          nchw_to_nhwc(flo_bwd),
                                          occ) * 100

            t_loss_feat = feature_temporal_loss(features,
                                                features_prev.detach(),
                                                nchw_to_nhwc(flo_bwd),
                                                occ) * 50
            
            #L1 loss for paired training
            if self.hparams.paired:
                l1_loss_AB = self.criterionL1(self.fake_B, self.real_B) * 100
                l1_loss_BA = self.criterionL1(self.fake_A, self.real_A) * 100
                l1_loss = (l1_loss_AB + l1_loss_BA) / 2

            ## DEBUG - show images while training

            if self.hparams.show_debug_im:
                plt.imshow(self.fake_B[0].detach().cpu().permute(1, 2, 0))
                plt.show()

            if self.hparams.paired:
                g_loss = l1_loss + t_loss + t_loss_feat  + loss_cycle_ABA + loss_cycle_BAB  + loss_GAN_A2B + loss_GAN_B2A
            else:
                g_loss = loss_GAN_A2B + loss_GAN_B2A + t_loss + t_loss_feat  + loss_cycle_ABA + loss_cycle_BAB 

            tqdm_dict = {'g_loss': g_loss, 
                         'GAN_loss_AB' : loss_GAN_A2B,
                         'GAN_loss_BA' : loss_GAN_B2A,
                         'temporal_loss' : t_loss, 
                         'temporal_feature_loss' : t_loss_feat, 
                         'cycle_loss_ABA' : loss_cycle_ABA, 
                         'cycle_loss_BAB' : loss_cycle_BAB                    
                         }
            
            if self.hparams.paired:
                tqdm_dict['L1_loss'] = l1_loss 

            self.last_imgs = self.fake_B
            
            self.log('train_loss', g_loss, prog_bar=True)

# I guess the problem must be here:

            output = OrderedDict({
                'loss': g_loss,
                'progress_bar': tqdm_dict,
                'log': tqdm_dict
            })
            
            return output
            
        
        if optimizer_idx == 1:
        
            real_A = batch['A']
            real_B = batch['B']

            ###### Discriminator A ######

            # Real loss
            pred_real = self.netD_A(torch.cat([self.real_A,self.real_A], dim=1))
            loss_D_real = self.criterion_GAN(pred_real, True)

            # Fake loss
            self.fake_BA = self.fake_BA_buffer.push_and_pop(self.fake_BA) 
            pred_fake = self.netD_A(self.fake_BA.detach())
            loss_D_fake = self.criterion_GAN(pred_fake, False)

            # Total loss
            loss_D_A = (loss_D_real + loss_D_fake)*0.5

            ###### Discriminator B ######

            # Real loss
            pred_real = self.netD_B(torch.cat([self.real_B,self.real_B], dim=1))
            loss_D_real = self.criterion_GAN(pred_real, True)

            # Fake loss
            self.fake_AB = self.fake_AB_buffer.push_and_pop(self.fake_AB)
            pred_fake = self.netD_B(self.fake_AB.detach())
            loss_D_fake = self.criterion_GAN(pred_fake, False)

            # Total loss
            loss_D_B = (loss_D_real + loss_D_fake)*0.5  
            
            d_loss = loss_D_B + loss_D_A
            
            tqdm_dict = {'d_loss': d_loss, 
                         'loss_D_B' : loss_D_B,
                         'loss_D_A' : loss_D_A }
            
            output = OrderedDict({
                        'loss': d_loss,
                        'progress_bar': tqdm_dict,
                        'log': tqdm_dict
                        })
            

            return output
        

                

    def configure_optimizers(self):

        
        print("Params to learn:")
        g_params = []
        for name,param in self.netG_A2B.named_parameters():
            if param.requires_grad == True:
                g_params.append(param)
                print("\t",name)
        for name,param in self.netG_B2A.named_parameters():
            if param.requires_grad == True:
                g_params.append(param)
                print("\t",name)  
                
        
        #g_params = list(self.netG_A2B.parameters()) + list(self.netG_B2A.parameters())
        d_params = list(self.netD_A.parameters()) + list(self.netD_B.parameters())
        
        #opt_g = torch.optim.SGD(g_params, lr=self.hparams.lr_g, momentum=0.9)
        #opt_g = torch.optim.Adam(g_params, lr=self.hparams.lr_g)
        opt_g = Ranger(g_params, lr=self.hparams.lr_g, k=20)
        opt_d = Ranger(d_params, lr=self.hparams.lr_d, k=20)#lr=0.0005
        
        #g_sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt_g, T_max=908)
        #g_sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt_g, patience=1, verbose=True)
        
        #d_sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt_d, patience=1, verbose=True)
        
        return [opt_g, opt_d], []
    
    
    
    
    # Logging images:
    
    #def on_epoch_start(self):
        #dummy_input = torch.randn(1, 3, 256, 512).cuda().detach()
        #self.logger.experiment.add_graph(self.netG_A2B, dummy_input)
    
    def on_epoch_end(self):
        
#        trainer.checkpoint_callback.on_epoch_end()
               
        grid = torchvision.utils.make_grid(self.last_imgs).cpu().detach() #denorm removed
        self.logger.experiment.add_image(f'generated_images', grid, self.current_epoch)  
        

[carmocca]

The original code was developed under 0.4 version (I think?)

You can check by accessing torch.load("your_checkpoint.ckpt")["pytorch-lightning_version"].

And then check out the release notes for all subsequent versions. You can also try upgrading version by version to see what breaks.