I have been trying to train a DF-GAN for text-to-image generation. However, after training for a while, the losses become NaN and after that the model does not recover from it. I am using Mixed Precision Training to decrease the training time and increase the batch_size.
Currently, on a V100 GPU (on Google Cloud), each epoch takes about 3 mins with mixed precision enabled. On disabling mixed precision, each epoch takes around 10 mins. So there’s no other option but to enable the mixed precision training.
Before DF-GAN, I was training StackGAN model and experienced the same behavior. I am not understanding as to why this is occurring. I  looked up the documentation for mixed precision training for models with multiple optimisers and gradient penalties. I incorporated all the changes that the documentation recommended to do but I could not find any improvements.
I will display the loss curves and the images generated by model after training DF-GAN for a while with MP enabled.

Screenshot from 2021-09-13 07-45-311789×361 38 KB


Screenshot from 2021-09-13 07-45-501920×1048 256 KB

The pictures were not purplish until the NaNs started to occur.
Below is the training configs that I am using to train the model.
cudnn.benchmark = True

class TrainingConfig:

    gen_learning_rate = 0.0001
    disc_learning_rate = 0.0004
    epsilon = 1e-8
    betas=(0.00,0.9)
    max_epochs = 600
    num_workers = 6
    batch_size = 32
    drop_last=True
    shuffle = True
    pin_memory = True
    ckpt_dir = "./DF-GAN-v1/"
    gen_ckpt_path = None
    disc_ckpt_path = None
    verbose = True
    device = "cuda"
    logdir = "dfganv1"
    snap_shot = 20
    
    def __init__(self,**kwargs) -> None:
        for key,value in kwargs.items():
            setattr(self,key,value)

Below is the training loop code for training the model.
def run_epoch(split):
            is_train = split == "train"
            if is_train:
                generator.train()
                discriminator.train()
            else:
                generator.eval()
                discriminator.eval()

            data = self.train_dataset if is_train else self.test_dataset
            loader = DataLoader(dataset=data, batch_size=config.batch_size,
                                shuffle=config.shuffle,
                                pin_memory=config.pin_memory,
                                num_workers=config.num_workers,
                                drop_last=config.drop_last)
            lossesD, lossesD_real, lossesD_wrong, lossesD_fake, losses_kl, losses_gen = [], [], [], [], [], []

            pbar = tqdm(enumerate(loader),total=len(loader)) if is_train and config.verbose else enumerate(loader)
            for it, data in pbar:
                # place data on the correct device
                images,captions,caption_len,class_ids, keys = prepare_data(data)
                hidden = text_encoder.init_hidden(config.batch_size)
                # words_embs: batch_size x nef x seq_len
                # sent_emb: batch_size x nef   
                word_embeddings, sentence_embeddings = text_encoder(captions, caption_len, hidden)
                word_embeddings, sentence_embeddings = word_embeddings.detach(), sentence_embeddings.detach()

                images = images[0].to(self.device)
                noise = torch.randn(config.batch_size,100,device=self.device)
                
                disc_optimizer.zero_grad(set_to_none=True)
                gen_optimizer.zero_grad(set_to_none=True)
                
                with amp.autocast():
                    real_features = discriminator(images)
                    output = discriminator.COND_DNET(real_features, sentence_embeddings)
                
                errorD_real = F.relu(1.0-output).mean()

                with amp.autocast():
                    output = discriminator.COND_DNET(real_features[:(config.batch_size-1)],sentence_embeddings[1:(config.batch_size)])
               errorD_wrong = F.relu(1.0+output).mean()

                    #synthesis fake images
                with amp.autocast():
                    fake_images = generator(noise,sentence_embeddings)
                    fake_features = discriminator(fake_images.detach())
                    output = discriminator.COND_DNET(fake_features,sentence_embeddings)

                errorD_fake = F.relu(1.0+output).mean()

                errorD = errorD_real + (errorD_wrong + errorD_fake) * 0.5

                scaler.scale(errorD).backward()
                scaler.step(disc_optimizer)
                scaler.update()

                #MA-GP (gradient penalty)
                interpolated = (images.data).requires_grad_()
                sent_inter = (sentence_embeddings.data).requires_grad_()
                
                with amp.autocast():
                    features = discriminator(interpolated)
                    out = discriminator.COND_DNET(features,sent_inter)

                    grads = torch.autograd.grad(outputs=out,
                                            inputs=(interpolated,sent_inter),
                                            grad_outputs=torch.ones(out.size()).cuda(),
                                            retain_graph=True,
                                            create_graph=True,
                                            only_inputs=True)
                     
                    grad0 = grads[0].view(grads[0].size(0), -1)
                    grad1 = grads[1].view(grads[1].size(0), -1)
                    grad = torch.cat((grad0,grad1),dim=1)                        
                    grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))
                    d_loss_gp = torch.mean((grad_l2norm) ** 6)
                    d_loss = 2.0 * d_loss_gp
                    
                disc_optimizer.zero_grad(set_to_none=True)
                gen_optimizer.zero_grad(set_to_none=True)

                scaler.scale(d_loss).backward()
                scaler.step(disc_optimizer)

    
                ### update G network ###
                disc_optimizer.zero_grad(set_to_none=True)
                gen_optimizer.zero_grad(set_to_none=True)
                with amp.autocast():
                    features = discriminator(fake_images)
                    output = discriminator.COND_DNET(features,sentence_embeddings)
                
                errorG = - output.mean()

                scaler.scale(errorG).backward()
                scaler.step(gen_optimizer)

                scaler.update()

                lossesD.append(errorD.item())
                lossesD_real.append(errorD_real.item())
                lossesD_fake.append(errorD_fake.item())
                lossesD_wrong.append(errorD_wrong.item())
                losses_gen.append(errorG.item())

Please let me know if the MP sections of my training loop are correct. If not let me know what other options are there to prevent NaN from occuring.
The model used is taken from : GitHub - tobran/DF-GAN: Deep Fusion Generative Adversarial Networks for Text-to-Image Synthesis