OOM in GAN

[ncuxomun]

🐛 Bug

GPU memory quickly builds up in few epochs

## Please reproduce using the BoringModel
class model(pl.LightningModule):
    def __init__(self, d_dim=100, z_dim=64, pcs_dim=64, zd_dim=20, lr=0.0002):
        super().__init__()
        # self.save_hyperparameters()
        self.lr = lr
        self.z_dim = z_dim
        self.zm_dim = pcs_dim
        self.zd_dim = zd_dim
        self.b1 = 0.9
        self.b2 = 0.999
        self.alpha = 1.0 # weight for losses
        self.beta = 1.0 # weight for zm reg losses
        self.k_step = 0

        ### initializing networks
        img_shape = (1, H, W)
        self.generator = Generator(img_shape, self.z_dim, self.zm_dim)
        self.discriminator = Discriminator(img_shape, self.zm_dim)
        self.dae = DAE(d_dim, self.zd_dim)
        self.zmzdreg = ZmZdReg(self.zm_dim, self.zd_dim)

        # application of weights
        self.generator.apply(self.weights_init)
        self.discriminator.apply(self.weights_init)
        self.dae.apply(self.weights_init)

        # Important: This property activates manual optimization.
        self.automatic_optimization = False  # True - Auto // # False - Manual update

        # for sampling
        self.validation_z = None
        self.pcs_sample = None
        self.imgs_ref = None
        self.logging_dict = {"g_loss": [], "g_avd_loss": [], "g_z_loss": [],
                             "d_loss": [], "d_avd_loss": [], "d_z_loss": [],
                             "data_loss": [], "reg_loss": []
                             }

    def forward(self, noise, pcs):
        maps_hr_out = self.generator(noise, pcs)
        return maps_hr_out

    def training_step(self, batch, batch_idx, optimizer_idx):  
        imgs, pcs, data = batch

        # sample noise
        z = torch.randn(imgs.shape[0], self.z_dim)
        z = z.type_as(imgs)

        # valid and fake labels
        valid = torch.ones(imgs.size(0), 1).type_as(imgs)
        fake = torch.zeros(imgs.size(0), 1).type_as(imgs)

        self.validation_z = z
        # to ensure same set to be generated
        if self.k_step == 0:
            self.pcs_sample = pcs
            self.imgs_ref = imgs

        # optimizers, manual access
        g_opt, d_opt, zd_opt, r_opt = self.optimizers()
        # self.r_opt = r_opt

        ########## train generator ##########
        fake_img = self(z, pcs)
        valid_pred, pcs_pred = self.discriminator(fake_img, pcs)
        g_loss = self.alpha*self.adv_loss(valid_pred, valid) + self.beta*self.mse_loss(pcs_pred, pcs)
        self.log('g_loss', g_loss, prog_bar=True)

        # Update generator gradients
        g_opt.zero_grad()
        self.manual_backward(g_loss, retain_graph=True)
        # Update generator optimizer
        g_opt.step()

        # dictionary logging
        self.logging_dict['g_loss'].append(g_loss)
        self.logging_dict['g_avd_loss'].append(self.alpha*self.adv_loss(valid_pred, valid))
        self.logging_dict['g_z_loss'].append(self.beta*self.mse_loss(pcs_pred, pcs))
        ######################################

        ########## train discriminator ##########
        valid_pred, pcs_pred = self.discriminator(imgs, pcs)
        real_loss = self.alpha*self.adv_loss(valid_pred, valid) + self.beta*self.mse_loss(pcs_pred, pcs)

        fake_detach = self(z, pcs).detach()
        fake_valid, fake_pcs_pred = self.discriminator(fake_detach, pcs)
        fake_loss = self.alpha*self.adv_loss(fake_valid, fake) + self.beta*self.mse_loss(fake_pcs_pred, pcs)

        d_loss = (real_loss + fake_loss)
        self.log('d_loss', d_loss, prog_bar=True)

        # Update gradients
        d_opt.zero_grad()
        self.manual_backward(d_loss, retain_graph=True)
        # Update optimizer
        d_opt.step()

        # dictionary logging
        self.logging_dict['d_loss'].append(d_loss)
        self.logging_dict['d_avd_loss'].append(self.alpha*(self.adv_loss(valid_pred, valid) + self.adv_loss(fake_valid, fake)))
        self.logging_dict['d_z_loss'].append(self.beta*(self.mse_loss(pcs_pred, pcs) + self.mse_loss(fake_pcs_pred, pcs)))
        ######################################

• PyTorch Lightning Version (e.g., 1.3.0): 1.3.8
• PyTorch Version (e.g., 1.8) 1.9
• Python version:
• OS (e.g., Linux): Windows 10
• CUDA/cuDNN version: 11.1
• GPU models and configuration: 2070 max-q
• How you installed PyTorch (conda, pip, source): pip
• If compiling from source, the output of torch.__config__.show():
• Any other relevant information:

Additional context

I have 4 models in total with 4 respective optimizers. But by disabling them two of them, I found that the error was being caused by the GAN parts. I honestly have no idea what could be a problem here because I did implement manual optimization for WGANGP before and I had no memory issues.

[ncuxomun]

I might have solved the problem by changing to:
self.manual_backward(x_loss, retain_graph=False). But I am not entirely sure.
Is there a doc page where I can look up description for manual_backward operation?

[justusschock]

The issue most likely is that when you append it to your logging dict (why do you do this at all, when we provide the log function?) You don't detach anything. Meaning that first you have an increasing number of values there (so memory will increase anyways) but second you store the autograd graph for each value which is pretty large and not necessary. So changing self.logging_dict['some_name].append(some_tensor) to self.logging_dict['some_name].append(some_tensor.detach()) should help a lot. Your memory will still be growing but only a bit (the amount of memory it needs to store the tensor). Ideally you remove those calls completely and just rely on self.log :)

[ncuxomun]

Hi, thanks a lot for a quick response. Yeah, it is really not smart of me to initiate another logger when there is one. I was trying to recall why I did that, and I think I opted to create another dictionary because I wanted to plot them together like:

logging_dict = torch.load('logging_dict.pkl', map_location=torch.device('cpu')) <-- I believe I had to do this because the logs were stored on the GPU.

plt.figure(figsize=(6, 4))

plt.plot(torch.FloatTensor(logging_dict['g_loss']), label='Total loss')

plt.plot(torch.FloatTensor(logging_dict['g_avd_loss']), label='G adv loss')
plt.plot(torch.FloatTensor(logging_dict['d_avd_loss']), label='D adv loss')
plt.plot(torch.FloatTensor(logging_dict['g_z_loss']), label='z loss')
plt.xlabel('# Iterations')
plt.ylabel('Loss')
plt.title('GAN training')
plt.legend()
plt.tight_layout()
plt.show()

Personally, I'd prefer to have more than one loss progression in a single plot like in the example above, do you think I can do that via self.log('x', x) or should I use log_dict as in:
values = {'loss': loss, 'acc': acc, ..., 'metric_n': metric_n}
self.log_dict(values)
?

[tchaton]

Hey @ncuxomun,

Yes, I believe it is related to storing tensors on GPU.

        self.logging_dict = {"g_loss": [], "g_avd_loss": [], "g_z_loss": [],
                             "d_loss": [], "d_avd_loss": [], "d_z_loss": [],
                             "data_loss": [], "reg_loss": []
                             }

[ncuxomun]

Hi, @tchaton. Thanks a lot for a quick response. I'll definitely remove redundant logging but I'd still prefer to access the log and plot the losses in a single plot mentioned in my reply above. Do you think I can save logs via model.log['x'] in the of a training process?

[tchaton]

You can still access the logged values at any point using trainer.callback_metrics.

class MyModel(LightningModule):

    def training_epoch_end(self, outputs):
        self.trainer.callback_metrics

[ncuxomun]

Thanks a lot, Thomas!

[ncuxomun]

@tchaton, I've a question regarding where to set backward pass and subsequently update the weights. Specifically, I have four sets of model defined in a single graph, which should have 3-4 individual optimizers. But I am not sure whether I am ought to manually update in the training_step or validation_step. Can you please have a look at my implementation and tell if I am doing something wrong here?

def training_step(self, batch, batch_idx, optimizer_idx):  
    imgs, pcs, data = batch

        # # optimizers, manual access
    g_opt, d_opt, zd_opt, r_opt = self.optimizers()

    # for model 1
    ########## train generator ##########
    # Update generator gradients
    g_opt.zero_grad()
    self.manual_backward(g_loss, retain_graph=True)
    # Update generator optimizer
    g_opt.step()

    ######################################

    ########## train discriminator ##########
   # for model 2
  # Update gradients
    d_opt.zero_grad()
    self.manual_backward(d_loss, retain_graph=True)
    # Update optimizer
    d_opt.step()
    ######################################
     
   # for model 3
    ########## train regressor ##########
    # Update gradients
    zd_opt.zero_grad()
    self.manual_backward(train_loss_data, retain_graph=True)
    # Update optimizer
    zd_opt.step()

    # for model 4
    # Update gradients
    r_opt.zero_grad()
    self.manual_backward(train_loss_zd, retain_graph=True)
    # Update optimizer
    r_opt.step()

def validation_Step(self, batch, batch_idx, optimizer_idx):  
    imgs, pcs, data = batch
   # omitted for brevity

OR

def training_step(self, batch, batch_idx, optimizer_idx):  
    imgs, pcs, data = batch
   # omitted for brevity

def validation_Step(self, batch, batch_idx, optimizer_idx): 
    # # optimizers, manual access
    g_opt, d_opt, zd_opt, r_opt = self.optimizers()

    # for model 1
    ########## train generator ##########
    # Update generator gradients
    g_opt.zero_grad()
    self.manual_backward(g_loss, retain_graph=True)
    # Update generator optimizer
    g_opt.step()

    ######################################

    ########## train discriminator ##########
   # for model 2
  # Update gradients
    d_opt.zero_grad()
    self.manual_backward(d_loss, retain_graph=True)
    # Update optimizer
    d_opt.step()
    ######################################
     
   # for model 3
    ########## train regressor ##########
    # Update gradients
    zd_opt.zero_grad()
    self.manual_backward(train_loss_data, retain_graph=True)
    # Update optimizer
    zd_opt.step()

    # for model 4
    # Update gradients
    r_opt.zero_grad()
    self.manual_backward(train_loss_zd, retain_graph=True)
    # Update optimizer
    r_opt.step()

So far, I am able to train with the former case, whilst the latter one does not work.

Thanks a lot for your attention!

[tchaton]

The optimization should be done in the training_step as you are training :) Looks great ! Exactly the right way to use manual optimization :) Neat !