OutOfMemoryError: CUDA out of memory

OutOfMemoryError: CUDA out of memory. Tried to allocate 1.87 GiB. GPU 0 has a total capacity of 21.66 GiB of which 1.85 GiB is free. Process 17782 has 210.00 MiB memory in use. Process 20882 has 272.00 MiB memory in use. Process 30219 has 224.00 MiB memory in use. Process 31600 has 1.66 GiB memory in use. Including non-PyTorch memory, this process has 17.46 GiB memory in use. Of the allocated memory 14.64 GiB is allocated by PyTorch, and 2.63 GiB 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)

遇到 OutOfMemoryError: CUDA out of memory 错误通常表示 GPU 内存不足以满足当前操作的需求。具体来说，PyTorch 尝试分配 1.87 GiB 的 GPU 内存，但当前 GPU 的剩余内存只有 1.85 GiB，这就导致了内存不足的错误。

原因分析

从错误信息中可以看出，PyTorch 已经分配了 14.64 GiB 的内存，而且还有 2.63 GiB 是被保留但未分配的。这样的内存碎片可能导致即使有 1.85 GiB 的空闲内存，仍然无法成功分配 1.87 GiB 的新内存。

解决方法

以下是一些可以尝试的解决方案：

1. 减少批处理大小 (Batch Size)

减少批处理大小是最常见的解决办法。批处理大小越大，需要的 GPU 内存就越多。如果你的模型在训练时遇到内存不足，可以尝试减小 batch_size。

例如：

batch_size = 16  # 原来的批处理大小，尝试减小到 8 或更小

2. 清理缓存

PyTorch 默认会缓存 GPU 内存，以提高性能，但有时内存碎片会导致分配失败。你可以手动清理缓存来释放内存：

import torch
torch.cuda.empty_cache()

调用 empty_cache() 后，PyTorch 会清理缓存，但不会释放已分配给变量的内存。

3. 检查和优化模型

• 减少模型的参数量：如果可能，减少网络层数或每一层的神经元数量，尤其是在处理大模型时。
• 使用 float16 精度：使用半精度浮动（16-bit 浮动）代替 32-bit 浮动，可以显著减少内存占用。PyTorch 支持混合精度训练，通常使用 torch.cuda.amp 自动混合精度（Automatic Mixed Precision）来减少内存使用。

# 使用自动混合精度（AMP）来减少内存消耗
from torch.cuda.amp import autocast

with autocast():
    # 在此块中运行前向和反向计算
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()

4. 使用 torch.utils.checkpoint

PyTorch 提供了一个内存优化的技巧——模型检查点（checkpointing）。它通过保存部分计算图，减少中间激活值的内存开销。在训练时，你可以选择将某些层的计算结果不存储，改为在反向传播时重新计算。

import torch.utils.checkpoint as checkpoint

# 对某些层使用检查点
output = checkpoint.checkpoint(model_layer, input)

5. 调整 PyTorch 的 CUDA 内存管理

根据错误信息，PyTorch 提供了 PYTORCH_CUDA_ALLOC_CONF 环境变量来控制内存分配策略。你可以尝试设置这个环境变量来避免内存碎片。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

这个设置允许 PyTorch 在内存碎片问题严重时扩展内存段，从而避免内存分配失败。

6. 重启Jupyter Notebook内核（如使用）

如果使用的是Jupyter Notebook的话， 可以尝试重启内核释放显存， 如果重启后显存仍然不足，尝试改小batch_size。

7. 尝试其他进程释放 GPU 内存

错误信息显示了其他进程在使用 GPU 内存，可能有其他进程占用了大量内存。你可以尝试结束不需要的进程来释放更多 GPU 内存。你可以使用如下命令查看 GPU 上的进程：

nvidia-smi

然后根据进程的 PID，可以选择结束某些不必要的进程：

kill -9 <pid>

8. 增加 GPU 内存

如果可能，尝试使用内存更大的 GPU。对于一些计算密集型任务，可能需要使用更强大的 GPU 来避免内存不足的错误。

9. 切换到多 GPU 设置

如果你的机器上有多个 GPU，可以尝试将计算任务分配到多个 GPU 上，分散内存负载。你可以使用 DataParallel 或 DistributedDataParallel 来在多个 GPU 上分配模型和数据。

model = torch.nn.DataParallel(model)
model = model.cuda()

总结

CUDA out of memory 错误通常是由于 GPU 内存不足或内存碎片造成的。你可以通过以下方式解决该问题：

1. 减小批处理大小。
2. 清理缓存，使用 torch.cuda.empty_cache()。
3. 优化模型，减少参数量。
4. 使用混合精度训练。
5. 使用内存优化技术，如模型检查点。
6. 设置 PYTORCH_CUDA_ALLOC_CONF 环境变量以减少内存碎片。
7. 结束其他进程，释放 GPU 内存。
8. 尝试多 GPU 计算或更大内存的 GPU。