PyTorch Lightning模型训练首个epoch后显存溢出如何优化

核心优化方案 #

单折内显存溢出修复（适配你第一折跑不完的场景） #

• 修复日志指标的显存占用
  这是PyTorch Lightning比原生PyTorch额外耗显存的最常见原因：PL默认会将你在training_step、validation_step中log的张量保留在GPU上，不会自动转移到CPU。你需要在Trainer初始化时添加move_metrics_to_cpu=True参数，同时在模型代码中所有log损失、指标的位置，都用loss.item()而非直接传入GPU上的loss张量，避免整个计算图被留存无法释放。
• 关闭PL默认的冗余功能
  你当前的Trainer没有关闭不必要的默认配置，都会额外占用显存：
  • 不需要打印模型结构的话添加enable_model_summary=False
  • 不需要默认 checkpoint 保存的话添加enable_checkpointing=False，如果需要保存模型仅保留最优的即可，配置save_top_k=1不要存所有epoch的权重
  • 旧版本PL可设置progress_bar_refresh_rate=10，避免进度条每步刷新产生冗余缓存
• 排查冗余张量移动
  PyTorch Lightning会自动把DataLoader返回的batch移动到GPU上，如果你在Classifier的training_step、forward方法中手动再次调用.to('cuda')转移张量，会生成冗余的GPU张量副本，直接删除手动移CUDA的代码即可。
• 避免计算图留存
  不要在类属性（比如self.loss_list）中直接存储GPU上的损失、预测张量，存储前必须调用.item()（单值）或者.detach().cpu().numpy()（数组），切断和计算图的关联。
• 排查验证步缓存问题
  旧版本PL会默认把所有validation_step的返回值存在GPU显存的列表里，直到validation_epoch_end运行完才释放，如果验证集规模大，第一次运行验证的时候就会爆显存。不需要的计算结果不要在validation_step中返回，必须返回的话提前转移到CPU。

多折训练显存泄漏修复 #

每折运行完成后手动清理残留显存，在run_fold函数末尾添加如下代码：

import gc
import torch
del model, trainer, train_dl, valid_dl, logger
gc.collect()
torch.cuda.empty_cache()

修改后Trainer参考代码 #

trainer = pl.Trainer(
    gpus = 1, 
    max_epochs = config.epochs,
    fast_dev_run = config.debug,
    logger = logger,
    log_every_n_steps=10,
    # 新增显存优化参数
    move_metrics_to_cpu=True,
    enable_model_summary=False,
    enable_checkpointing=False # 需要保存模型的话自行调整对应参数
)

内容的提问来源于stack exchange，提问作者Abhishek Prajapat