PyTorch 显存占用怎么排查？看已分配、缓存与峰值

AI小管家 • 7小时前 • AI教程 • 阅读 0

看到显存占用很高，先区分两件事：当前张量实际占用多少内存，以及 PyTorch 缓存分配器保留了多少内存。两者不同，直接反复调用 empty_cache() 不会释放仍被张量持有的显存。下面用 PyTorch 自带计数器定位峰值，依据 PyTorch 2.14 CUDA 内存管理说明。

在同一进程的关键位置打点

先确认 CUDA 可用。在已有训练或推理脚本里加入：

import torch

if not torch.cuda.is_available():
    raise RuntimeError("当前环境没有可用 CUDA 设备")

def report(stage):
    mib = 1024 ** 2
    torch.cuda.synchronize()
    print(stage,
          "allocated MiB", round(torch.cuda.memory_allocated() / mib, 1),
          "reserved MiB", round(torch.cuda.memory_reserved() / mib, 1),
          "peak allocated MiB", round(torch.cuda.max_memory_allocated() / mib, 1))

report("模型加载前")
model = torch.nn.Linear(256, 256).eval().cuda()
report("模型加载后")
x = torch.randn(32, 256, device="cuda")
with torch.inference_mode():
    model(x)
report("一次推理后")

小型线性层和随机输入只是演示打点位置，没有代表任何实际模型。把 report() 放到自己已有脚本的模型加载前后、输入搬到 GPU 后和一次完整训练步后，再用同一批输入重跑。max_memory_allocated() 是当前进程迄今记录的张量分配峰值；独立启动一次脚本便于比较某项改动前后的峰值。

怎样读差值并尝试调整

如果模型加载后 allocated 明显上升，先核对模型权重和额外模块；若训练步后峰值跳升，检查批大小、输入长度、中间激活和是否无意保留了每步的计算图。一次只改一个因素，例如减小批大小，再比较同一阶段的峰值及任务输出是否仍正确。

官方说明：memory_allocated() 反映张量占用，memory_reserved() 反映缓存分配器管理的总量；empty_cache() 只能释放未使用的缓存，不能释放仍由张量占用的内存。若仍无法解释，可以按官方 CUDA 内存快照文档进一步记录分配历史，但快照本身也有开销。

验证方法是同一模型与输入在调整前后都能完成目标步骤，并记录实际峰值。若发生 CUDA OOM、计数器与外部工具不一致，先核对是否有其他进程或非 PyTorch 分配；官方也指出 PyTorch 的内存分析不覆盖所有第三方 CUDA 分配。本文没有对具体 GPU、模型或量化格式作“必能运行”的判断。

Ai菜鸟网。发布者：AI小管家，转载请注明出处：https://www.alyyhw.com/29928.html