PyTorch CUDA out of memory报错的代码修改解决方法

你的GPU总显存为10.76GiB，原代码存在无意义的设备拷贝、显存未及时释放的问题，哪怕把batch size降到极低也容易触发显存溢出，按以下步骤修改即可解决：

• 先清理残留显存
  执行nvidia-smi查看当前GPU占用进程，找到残留的Python、训练相关进程，用kill -9 进程号杀掉，释放被无效占用的显存。日志里的Gdk-CRITICAL是无图形界面服务器的GTK组件警告，和显存问题无关，直接忽略即可。
• 修复PSM类的显存冗余写法
  原forward方法里反复调用.cuda()做CPU-GPU张量拷贝，会导致PyTorch反复申请显存不释放，是报错的核心原因，替换为以下代码：

  def forward(self, x):
      encodes = []
      outputs = []
      for device_data in x:
          # 输入数据统一搬到GPU，encoder输出直接留在显存中，避免来回拷贝
          encode = self.encoder(device_data.cuda())
          outputs.append(self.decoder(encode))
          encodes.append(encode)
      # 共享特征计算全程在GPU完成，无额外拷贝开销
      shared_encode = torch.mean(torch.stack(encodes), dim=2).permute(1,0,2)
      outputs.append(self.decoder(shared_encode))
      # 取完结果立刻删除无用中间张量，主动回收显存
      final_output = torch.mean(torch.stack(outputs), dim=0)
      del encodes, outputs, shared_encode
      torch.cuda.empty_cache()
      return final_output
  

• 训练侧显存优化配置
  • 开启混合精度训练，可降低40%左右显存占用，对精度几乎无影响：
    1. 训练初始化阶段添加scaler = torch.cuda.amp.GradScaler()
    2. 前向传播逻辑包裹在with torch.cuda.amp.autocast():上下文内
    3. 反向传播替换为scaler.scale(loss).backward()，参数更新用scaler.step(optimizer)搭配scaler.update()
  • 初始化模型时将num_feat_map参数从默认64调整为32，该参数控制卷积通道数，直接决定模型参数量和显存占用，减半后精度损失极小，显存占用可降低近50%
  • 训练循环中记录loss时取纯数值，用loss.item()存储，不要直接保留带计算图的tensor对象，避免计算图长期占用显存
  • 每个batch训练结束后调用一次torch.cuda.empty_cache()清空显存缓存，不需要额外加冗余的全局垃圾回收逻辑
• 极端小显存场景补充优化
  如果以上修改后仍有显存压力，开启梯度检查点功能，以20%左右的训练速度损耗换取30%的显存节省：模型初始化后，对encoder、decoder的卷积层调用torch.utils.checkpoint.checkpoint()包装即可。

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