training_step with autocast(endabled=True) and GradScaler()

[kapsner]

Hi,

I would like to reimplement this code with lightning and I am not sure, how to correctly write the training_step.

I've implemented something like the following but I am unsure, if this is the correct way of transposing the code to the lightning framework:

def __init__(self, **kwargs):
    self.scaler_g = torch.cuda.amp.GradScaler()
    self.scaler_d = torch.cuda.amp.GradScaler()
    self.automatic_optimization = False

def training_step(self, batch, batch_idx, prefix):

        images = batch["image"]
        # Get optimizers here
        optimizer_g, optimizer_d = self.optimizers()

        # Generator part
        if prefix == "train":
            optimizer_g.zero_grad(set_to_none=True)

        with autocast(enabled=True):
            reconstruction, z_mu, z_sigma = self.autoencoder(images)
            kl_loss = self.KL_loss(z_mu, z_sigma)

            recons_loss = self.l1_loss(input=reconstruction.float(), target=images.float())
            p_loss = self.loss_perceptual(
                reconstruction.float(),
                images.float()
            )
            loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss

            if self.current_epoch > self.hparams.autoencoder_warm_up_n_epochs:
                logits_fake = self.discriminator(reconstruction.contiguous().float())[-1]
                generator_loss = self.adv_loss(logits_fake, target_is_real=True, for_discriminator=False)
                loss_g += adv_weight * generator_loss
        
        if prefix == "train":
            self.scaler_g.scale(loss_g).backward()
            self.scaler_g.step(optimizer_g)
            self.scaler_g.update()

        if self.current_epoch > self.hparams.autoencoder_warm_up_n_epochs:
            # Discriminator part
            if prefix == "train":
                optimizer_d.zero_grad(set_to_none=True)

            with autocast(enabled=True):
                logits_fake = self.discriminator(reconstruction.contiguous().detach())[-1]
                loss_d_fake = self.adv_loss(logits_fake, target_is_real=False, for_discriminator=True)
                logits_real = self.discriminator(images.contiguous().detach())[-1]
                loss_d_real = self.adv_loss(logits_real, target_is_real=True, for_discriminator=True)
                discriminator_loss = (loss_d_fake + loss_d_real) * 0.5

                loss_d = adv_weight * discriminator_loss
            
            if prefix == "train":
                self.scaler_d.scale(loss_d).backward()
                self.scaler_d.step(optimizer_d)
                self.scaler_d.update()
        
        return_dict = {"recons_loss": recons_loss.item()}

        if self.current_epoch > self.hparams.autoencoder_warm_up_n_epochs:
            return_dict = dict(return_dict, **{"gen_loss": generator_loss.item()})
            return_dict = dict(return_dict, **{"disc_loss": discriminator_loss.item()})
        
        return return_dict

In particular, I am unsure if it would be sufficient to remove the GradScaler() code and autocast() and replace it with trainer = Trainer(precision=16) (as suggested here), so that the code would look something like:

def __init__(self, **kwargs):
    self.automatic_optimization = False

def training_step(self, batch, batch_idx, prefix):

        images = batch["image"]
        # Get optimizers here
        optimizer_g, optimizer_d = self.optimizers()

        # Generator part
        if prefix == "train":
            optimizer_g.zero_grad(set_to_none=True)

        reconstruction, z_mu, z_sigma = self.autoencoder(images)
        kl_loss = self.KL_loss(z_mu, z_sigma)

        recons_loss = self.l1_loss(input=reconstruction.float(), target=images.float())
        p_loss = self.loss_perceptual(
            reconstruction.float(),
            images.float()
        )
        loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss

        if self.current_epoch > self.hparams.autoencoder_warm_up_n_epochs:
            logits_fake = self.discriminator(reconstruction.contiguous().float())[-1]
            generator_loss = self.adv_loss(logits_fake, target_is_real=True, for_discriminator=False)
            loss_g += adv_weight * generator_loss
        
        if prefix == "train":
            self.manual_backward(loss_g)
            optimizer_g.step()

        if self.current_epoch > self.hparams.autoencoder_warm_up_n_epochs:
            # Discriminator part
            if prefix == "train":
                optimizer_d.zero_grad(set_to_none=True)

            logits_fake = self.discriminator(reconstruction.contiguous().detach())[-1]
            loss_d_fake = self.adv_loss(logits_fake, target_is_real=False, for_discriminator=True)
            logits_real = self.discriminator(images.contiguous().detach())[-1]
            loss_d_real = self.adv_loss(logits_real, target_is_real=True, for_discriminator=True)
            discriminator_loss = (loss_d_fake + loss_d_real) * 0.5

            loss_d = adv_weight * discriminator_loss
            
            if prefix == "train":
                self.manual_backward(loss_d)
                optimizer_d.step()
        
        return_dict = {"recons_loss": recons_loss.item()}

        if self.current_epoch > self.hparams.autoencoder_warm_up_n_epochs:
            return_dict = dict(return_dict, **{"gen_loss": generator_loss.item()})
            return_dict = dict(return_dict, **{"disc_loss": discriminator_loss.item()})
        
        return return_dict