如何限制PyTorch的GPU显存预留量？已尝试环境变量仍报错

报错信息 #

CUDA out of memory. Tried to allocate 2.52 GiB (GPU 0; 6.00 GiB total capacity; 4.59 GiB already allocated; 0 bytes free; 4.61 GiB reserved in total by PyTorch)

通用解决方案 #

1. 用PyTorch原生API限制显存 #

• 按比例限制进程显存：通过torch.cuda.set_per_process_memory_fraction()指定当前进程可使用的GPU显存占总容量的比例，比如限制为70%：

  import torch
  torch.cuda.set_per_process_memory_fraction(0.7, device=0)  # device=0对应第1块GPU
  

• 启用动态显存分配：设置PYTORCH_CUDA_ALLOC_CONF环境变量，让显存按需分配，避免一次性预留过多：

  # Linux/macOS终端设置
  export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  # Windows命令行设置
  set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  

  也可以在Python代码中直接设置：

  import os
  os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
  

2. 限制卷积操作的Workspace大小（替代调试参数） #

CUDNN_CONV_WSCAP_DBG是调试用参数，建议用官方支持的方法控制卷积显存：

import torch
from torch.backends import cudnn

# 设置最大workspace大小，单位为字节（示例为2GB）
cudnn.set_max_workspace_size(2 * 1024**3)

这个方法和YOLO8的workspace参数原理一致，属于通用的卷积显存控制方式。

3. 辅助显存优化手段 #

• 梯度累积：将多个batch的梯度累积后再更新参数，降低单次显存占用：

  accum_steps = 4  # 累积4个batch的梯度再更新
  for i, (inputs, labels) in enumerate(dataloader):
      outputs = model(inputs)
      loss = criterion(outputs, labels)
      loss = loss / accum_steps  # 损失按累积步数均分
      loss.backward()
      if (i + 1) % accum_steps == 0:
          optimizer.step()
          optimizer.zero_grad()
  

• 模型量化：将模型从FP32转为INT8，大幅压缩显存占用：

  model = model.to('cuda')
  model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
  torch.quantization.prepare(model, inplace=True)
  # 用少量数据校准模型
  for inputs, _ in dataloader:
      model(inputs.to('cuda'))
      break
  torch.quantization.convert(model, inplace=True)
  

• 主动释放缓存：清理无用张量后调用empty_cache()释放显存：

  del unused_tensor  # 删除不再使用的张量
  torch.cuda.empty_cache()  # 释放显存缓存
  

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