Why does pytorch lightning cause more GPU memory usage?

[chuzheng88]

Assumign that my model uses 2G GPU memory, every batch data uses 3G GPU memory. Traning code will use 5G (2+3) GPU memory when I use Pytorch. However, new training code use 8G (2 + 3 + 3) GPU memory when I use Pytorch ligntning. I do not know why ? or new codes are wrong ?
Some new codes described as follows:

Can you help me ?

[akihironitta]

Hi @chuzheng88! Thank you for waiting for replies patiently here. I haven't looked at your pure PyTorch implementation, so I can't point out issues for sure. However, to me, the suspect looks like the return value of traininig_step. PL accumulates its return values across steps in an entire epoch so that they can be passed to training_epoch_end as the input. I'd suggest not returning "pred_y", "y", "idx" if you don't plan to use their values there.
https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#training-epoch-end

Also, there is a list of possible causes in performance difference between PL and PyTorch in #12398.

[chuzheng88]

Hi @akihironitta , thank you for your reply. According to your suggestion, I can rmove the function "training_epoch_end()". But I need to log epoch's level error for saving the best model, so what should I do ?

[akihironitta]

If you're using one of the metrics supported by torchmetrics, I'd recommend using one from the library. It aggregates metric values in a single value instead of holding them all in memory across batches and aggregating them after each epoch.

Some doc pages that you might find helpful are:

• TorchMetrics with PL: https://torchmetrics.readthedocs.io/en/stable/pages/lightning.html
• Note on differentiability: https://torchmetrics.readthedocs.io/en/stable/pages/overview.html#metrics-and-differentiability
• MSE loss: https://torchmetrics.readthedocs.io/en/stable/regression/mean_squared_error.html

Let me know if it works for your case, and feel free to ask me any further questions here, in another discussion, or in DM on PyTorch Lightning Slack! :)

[chuzheng88]

@akihironitta Thank you for your reply, I wil try the TorchMetrics, and then report on the result.

[chuzheng88]

Hi @akihironitta . According to your suggestion, I read the official exampl as follow:

And then, I modified my code as follows:
class BPRegression(pl.LightningModule):
def init(self, model, train_dataset, val_dataset, test_dataset):
super().init()

    self.model = model
    self.train_loss = torchmetrics.MeanSquaredError()
    self.val_loss = torchmetrics.MeanSquaredError()
    self.lr = config_params["lr"]
    self.model_train_loss = 0
    self.model_val_loss = 0        

def _cal_loss(self, pred, y):
    loss = F.mse_loss(pred, y)
    return loss

def training_step(self, batch, batch_idx):
    idx, y, p_info, ppg_filter_1000 = batch
    
    # file_info = torch.squeeze(file_info)
    idx, p_info, ppg_filter_1000 = idx.reshape(-1), p_info.reshape(-1, 5), ppg_filter_1000.reshape(-1, 1, 2717)
    pred = self.model(ppg_filter_1000, p_info)
    pred, y = pred.squeeze(1), y.squeeze()
    
    # # opt = self.optimizers()
    # # opt.zero_grad()
    # loss = self._cal_loss(pred, y)
    # # self.manual_backward(loss)
    # # opt.step()
    
    self.train_loss(pred, y)
    self.log(f"loss/train_loss", self.train_loss, on_step=False, on_epoch=True)
    # return {"loss": loss}

# def training_epoch_end(self, outputs):
#     self.model_train_loss = torch.mean(torch.tensor([o["loss"] for o in outputs]))
#     self.log('loss/train_loss', self.model_train_loss)
#     # plot train and val loss in one picture, note that it will cause multi-folder problem
#     self.logger.experiment.add_scalars(f"Loss_on_epoch", {"train":self.model_train_loss, "val":self.model_val_loss}, self.current_epoch)
#     self.log_dict({f"test_epoch_train_loss":self.model_train_loss,f"test_epoch_val_loss":self.model_val_loss}, on_step=True, on_epoch=True)
#     # plot model weights and gradients
#     for name, parameters in self.model.named_parameters():
#         self.logger.experiment.add_histogram(f"Model/{name}", parameters, self.current_epoch)
#     # self._cal_metrics("Train", outputs)
#     print("Epoch:", self.current_epoch, "train loss:", self.model_train_loss.item(), "val loss:", self.model_val_loss.item())
    
def validation_step(self, batch, batch_idx):
    idx, y, p_info, ppg_filter_1000 = batch
    # print(idx.shape, p_info.shape, ppg_filter_1000.shape)
    # file_info = torch.squeeze(file_info)
    idx, p_info, ppg_filter_1000 = idx.reshape(-1), p_info.reshape(-1, 5), ppg_filter_1000.reshape(-1, 1, 2717)
    pred = self.model(ppg_filter_1000, p_info)
    pred, y = pred.squeeze(1), y.squeeze()
    # loss = self._cal_loss(pred, y)
    self.val_loss(pred, y)
    self.log("loss/val_loss", self.val_loss, on_step=False, on_epoch=True)
    self.log("ckpt_val_loss", self.val_loss, on_step=False, on_epoch=True)
    # return {"loss": loss}

# def validation_epoch_end(self, outputs):
#     self.model_val_loss = torch.mean(torch.tensor([o["loss"] for o in outputs]))
#     self.log("loss/val_loss", self.model_val_loss)
#     # log another val loss for checkpoint monitor
#     self.log("ckpt_val_loss", self.model_val_loss)
#     # self._cal_metrics("Val", outputs)

I find that the modified code was not work when computing MSE error.

In my opinion, the modified code not do backpropagation operation accroding to MSE error.

In another official example as follows:

This method needs the function "training_epoch_end()", It may be cause higher GPU memory usage.

So, what can I do ?

[chuzheng88]

OK, I have solved this problem after returning loss on the function "training_step()". I will test the GPU memory usage.

Here, another problem is the funciton "self.log(f"loss/train_loss", self.train_loss, on_step=False, on_epoch=True)", the X-axis shows Steps.