第3章 显存监控工具：nvidia-smi、torch.cuda.memory_summary 等

显存优化这事儿，说白了就跟修车一样。你得先知道哪儿漏了，才能动手去补。

我刚入行那会儿，靠直觉调批量大小。结果模型要么跑着跑着崩了，要么显存还剩一大半，性能没榨干。后来才明白——监控是第一步。工具就那几个，但用法够你琢磨一阵的。

3.1 最常用的就是 nvidia-smi

这是 NVIDIA 官方命令行工具。你装了驱动就能用。我习惯在训练脚本里隔几秒调一次，记录显存峰值。
嗯，这里要注意，它查的是 全局显存占用，不是单进程的。

# 实时刷新（每秒一次）
nvidia-smi -l 1

# 只看某个 GPU (比如 0 号)
nvidia-smi -i 0

# 查询指定指标，输出给脚本用
nvidia-smi --query-gpu=index,timestamp,memory.used,memory.total,utilization.gpu --format=csv -l 1

用 --query-gpu 是我最喜欢的。你可以把它写进日志，训练完了回看显存曲线。我曾经有个项目，模型在某个 batch 突然爆 OOM，就是靠 csv 日志抓到的时间点——原来是在数据加载时多了一倍显存开销。

3.2 PyTorch 自带的内存快照与摘要

光靠 nvidia-smi 不够。它只能告诉你总量，但到底哪块代码在吃显存？是模型参数、梯度、还是中间激活？
这时候就得请出 PyTorch 的内部监控了。

3.2.1 memory_summary

一个人就能调出来的详细报告。个人建议放到每个 epoch 结尾，或者 OOM 异常捕捉里。

import torch

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
data = torch.randn(4, 3, 224, 224).cuda()
out = model(data)
loss = out.sum()
loss.backward()
optimizer.step()

print(torch.cuda.memory_summary())

它会打印一堆东西。你重点看这几行：

• Allocated memory —— 当前真正在用的显存
• Active memory —— 正在被计算使用的
• Reserved memory —— PyTorch 向驱动预占的，不一定全用到
• Large alloc pool 和 Small alloc pool —— 我经常从这儿看出有没有碎片化

你想想看，如果 reserved 比 allocated 大很多，说明缓存策略没调好。这时候调 torch.cuda.set_per_process_memory_fraction 或者 max_split_size_mb 就能回收不少。

3.2.2 memory_snapshot + Chrome Trace

要是 memory_summary 还不够细，就上 snapshot。它会记录每个张量的分配栈，你可以在 Chrome 里看火焰图。

# 记录快照
torch.cuda.memory_snapshot()
# 导出为 Chrome trace 格式
from torch.cuda import _memory_trace
trace = _memory_trace._dump_snapshot(...)

# 或者直接走新的 API（推荐）
with torch.profiler.profile(profile_memory=True) as prof:
    out = model(data)
    loss = out.sum()
    loss.backward()
prof.export_chrome_trace("trace.json")

为什么用 Chrome trace？
因为它的时间轴很清楚。你能看到每个算子启动时显存的变化，哪个峰值、哪个泄露，一目了然。

3.2.3 主动检测与告警

光看报告太被动了。我习惯在训练脚本里加两个检查点：

def check_memory(threshold_mb=7000):
    used = torch.cuda.memory_allocated() / 1024**2
    if used > threshold_mb:
        print(f"[WARNING] 显存占用超过 {threshold_mb} MB: {used:.1f} MB")
        # 或者直接中止本次 batch，避免 OOM 拖垮整个训练

# 每次迭代前检查
for batch in dataloader:
    check_memory(threshold_mb=7000)
    # forward / backward ...

3.3 知识点结构一览

上面说了这么多，其实核心就三层：

• 命令行级：nvidia-smi 快速看全局
• API 级：memory_summary、allocated/reserved 看本进程
• 主动诊断级：snapshot + 自定义阀值检查

3.4 对比一下：什么时候用什么

3.5 总结一下关键经验

• nvidia-smi 看的是全局，不是你自己的进程。多卡或多进程时别被它骗了。
• memory_summary 是诊断第一步，但别老开着，它本身也有开销。
• Snapshot + Trace 是终极武器，我一般在模型改结构、换框架版本后才用一次。
• 养成习惯：每次调整 batch size 或模型结构后，先跑一个短实验，检查显存曲线。