PyTorch下FHD视频训练CNN-LSTM出现显存溢出问题求助

• 计算图残留未释放：代码中直接打印loss张量而非标量值，会导致每轮迭代的损失张量持续持有完整计算图引用，中间激活无法被PyTorch自动回收，显存随迭代次数持续上涨。同时LSTM循环中初始化的隐藏态(lstm_hidden, lstm_cell)未显式detach，上一轮的计算图会和本轮耦合，进一步占用显存。
• CNN中间激活无优化：1080p分辨率的图像输入CNN后，中间层激活占用显存极高，未启用梯度检查点的情况下，所有训练阶段的中间激活都会完整保存在显存中，单batch 2个样本、共120帧的激活总占用远超过单卡11G显存阈值。
• 多GPU负载不均衡：从lstm_model.module调用可判断使用了nn.DataParallel封装模型，该实现默认会将所有batch的损失计算、梯度汇总都放在0号主卡上，即使总显存足够，单主卡显存占满也会触发OOM。
• 冗余显存占用：预先在GPU上初始化cnn_results、lstm_results两个大张量，且赋值过程中保留了所有CNN输出的梯度引用，额外占用了不必要的显存空间。

1. 修复计算图残留问题
   • 将print(loss)改为print(loss.item())，仅打印损失的标量值，释放计算图引用
   • LSTM隐藏态初始化后增加detach操作，切断与上一轮计算图的关联：

     (lstm_hidden, lstm_cell) = lstm_model.module.initLSTM()
     lstm_hidden = lstm_hidden.detach()
     lstm_cell = lstm_cell.detach()
     

   • 修正损失函数调用错误：函数参数传入的是loss_fn，但实际调用的是全局变量criterion，替换为统一变量避免潜在逻辑错误。
2. 增加显存优化配置
   • 为CNN模型启用梯度检查点，可降低70%左右的中间激活显存占用，只需在CNN模型定义后添加梯度检查点封装即可。
   • 调整张量存储逻辑，将cnn_results初始化放在CPU上仅做结果汇总，若仅用CNN做特征提取不需要训练，可在CNN前向逻辑外层加with torch.no_grad()进一步节省显存。
3. 优化多GPU分发策略
   弃用nn.DataParallel，改用nn.parallel.DistributedDataParallel做多卡训练，可实现显存和计算的均匀负载，避免单卡先爆显存的问题。
4. 其他可选优化
   • 先调低batch size验证训练逻辑是否正常，再逐步上调
   • 训练阶段开启自动混合精度训练，可再降低50%左右的显存占用：

     scaler = torch.cuda.amp.GradScaler()
     # 前向阶段用autocast封装
     with torch.cuda.amp.autocast():
         cnn_result = cnn_model(batch)
         # 其余前向、损失计算逻辑也放在该上下文内
     # 反向传播改用scaler
     scaler.scale(loss).backward()
     scaler.step(optimizer)
     scaler.update()
     

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