PyTorch训练出现CUDA Out of Memory问题，求有效解决方案

首先是报错信息：

torch.cuda.OutOfMemoryError: CUDA out of memory.

 Tried to allocate 3.52 GB. GPU 0 has a total capacity of 11.75 GB of which 609.19 MB is free. Including non-PyTorch memory, this process has 10.89 GB memory in use. Of the allocated memory 10.72 GB is allocated by PyTorch, and 17.72 MB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting 

    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)    

以下是按优先级排序的有效解决方案：

• 优化内存碎片（错误提示方案）
  启动训练脚本前设置环境变量，减少内存碎片：

  export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  

• 调小Batch Size
  直接降低训练的batch size，这是最快速见效的方法。若不想牺牲训练效果，可配合梯度累积使用。

• 启用混合精度训练
  利用torch.cuda.amp模块将部分张量转为半精度（FP16），大幅降低内存占用：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  for epoch in epochs:
      for inputs, labels in dataloader:
          optimizer.zero_grad()
          with autocast():
              outputs = model(inputs)
              loss = loss_fn(outputs, labels)
          scaler.scale(loss).backward()
          scaler.step(optimizer)
          scaler.update()
  

• 梯度累积
  累积多个小batch的梯度后再更新参数，等价于使用大batch size但不占用对应内存：

  accumulation_steps = 4  # 累积4个batch的梯度
  for epoch in epochs:
      for i, (inputs, labels) in enumerate(dataloader):
          optimizer.zero_grad()
          outputs = model(inputs)
          loss = loss_fn(outputs, labels)
          loss = loss / accumulation_steps  # 均分损失值
          loss.backward()
          if (i + 1) % accumulation_steps == 0:
              optimizer.step()
  

• 主动释放无用张量并清空缓存
  手动删除不再使用的张量，再清空CUDA缓存：

  import torch
  import gc
  
  # 删除无用张量
  del unused_tensor
  # 强制垃圾回收+清空CUDA缓存
  gc.collect()
  torch.cuda.empty_cache()
  

  建议在epoch结束后、验证阶段前执行，确保释放训练过程中产生的临时张量。

• 将模型转为半精度
  直接把模型参数转为FP16格式，进一步压缩内存：

  model = model.half()
  inputs = inputs.half()  # 输入张量需同步转半精度
  

  注意部分算子不支持FP16，遇到报错时需单独处理相关层。

• 检查内存泄漏
  使用torch.cuda.memory_summary()打印内存使用详情，定位持续占用内存的模块或张量。可在每个batch结束后打印，观察内存是否持续增长。

• 关闭不必要的梯度计算
  推理或验证阶段无需计算梯度，用torch.no_grad()上下文管理器禁用梯度：

  with torch.no_grad():
      val_outputs = model(val_inputs)
      val_loss = loss_fn(val_outputs, val_labels)
  

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