PyTorch持续出现CUDA内存不足问题求助

尽管已将batch size调整为8，PyTorch仍持续抛出CUDA显存不足错误。尝试过torch.cuda.empty_cache()清理显存，也试过相关显存优化方案，但均无效。

错误日志如下：

Traceback (most recent call last):
  File "D:\Programming\MachineLearning\Projects\diffusion_models\practice\ddpm.py", line 110, in <module>
    launch()
  File "D:\Programming\MachineLearning\Projects\diffusion_models\practice\ddpm.py", line 106, in launch
    train(args)
  File "D:\Programming\MachineLearning\Projects\diffusion_models\practice\ddpm.py", line 85, in train
    loss.backward()
  File "D:\Programming\global_venv\lib\site-packages\torch\_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "D:\Programming\global_venv\lib\site-packages\torch\autograd\__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
RuntimeError: CUDA out of memory. Tried to allocate 1024.00 MiB (GPU 0; 4.00 GiB total capacity; 2.63 GiB already allocated; 0 bytes free; 2.73 GiB reserved in total by PyTorch)

针对4G显存运行DDPM这类扩散模型的场景，以下方案能有效缓解显存压力：

• 进一步缩小batch size：直接降到4甚至2，扩散模型的隐层特征和噪声张量本身占用显存高，8对4G显存还是偏大。
• 启用梯度累积：如果不想太小的batch size影响训练稳定性，每N次前向传播后再执行一次反向传播。代码示例：

  accumulation_steps = 4  # 累计4次小batch等价于batch size 32
  for i, (data, _) in enumerate(dataloader):
      outputs = model(data)
      loss = criterion(outputs, targets)
      loss = loss / accumulation_steps  # 损失按步数均分
      loss.backward()
      if (i + 1) % accumulation_steps == 0:
          optimizer.step()
          optimizer.zero_grad()
  

• 使用混合精度训练：开启PyTorch的自动混合精度，能大幅降低显存占用。代码示例：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  for data in dataloader:
      optimizer.zero_grad()
      with autocast():
          outputs = model(data)
          loss = criterion(outputs, targets)
      scaler.scale(loss).backward()
      scaler.step(optimizer)
      scaler.update()
  

• 减少模型参数量：比如把DDPM的UNet编码器层数减少，或者降低特征图的通道数（比如从64改成32）。
• 清理计算图冗余：确保每次反向传播后及时清空梯度，避免不必要的张量保留，同时避免在训练循环中创建新的计算图节点。
• 关闭不必要的CUDA进程：用nvidia-smi查看GPU占用，关掉其他占用显存的程序（比如浏览器、其他PyTorch进程）。

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