4、诊断工具入门：nvidia-smi基础使用、torch.cuda API、内存快照工具

好，咱们正式开始动手了。

前面聊了那么多显存溢出的原理，你可能会想：“道理我都懂，可到底怎么定位问题？”

别急，这一章就是给你递工具的。我个人的习惯是，先学会看仪表盘，再学修发动机。显存诊断也一样，你得先知道“哪里满了”、“谁在吃显存”，然后才能对症下药。

4.1 nvidia-smi：GPU 界的“任务管理器”

说到 GPU 监控，nvidia-smi 绝对是第一把交椅。你想想看，每次训练崩了，第一反应是不是敲一下这个命令？

基础用法很简单，终端里直接输入：

nvidia-smi

你会看到一张表格，里面列出了每张卡的：

• 温度（Temp）：太高了会降频，影响训练速度。
• 显存使用（Memory-Usage）：比如 1234MiB / 24576MiB，这就是核心指标。
• GPU 利用率（GPU-Util）：注意，这个不是显存占用率，是计算单元忙不忙。
• 进程列表（Processes）：哪个进程在用这张卡，占了多少显存。

但 nvidia-smi 有个坑——它默认只刷新一次。我建议你养成习惯，用动态监控模式：

nvidia-smi -l 2

这个 -l 2 表示每 2 秒刷新一次。嗯，这里要注意，刷新频率别设太低，比如 0.1 秒，否则终端刷得你眼花，而且会占用一点点 GPU 资源去查询。

我曾经在调试一个 8 卡分布式训练时，发现其中一张卡的显存莫名其妙比别人高 2GB。用 nvidia-smi -l 1 盯着看了半分钟，才发现是数据加载进程的 num_workers 设置不一致导致的。这种“肉眼盯梢”的方法，有时候比代码调试还管用。

4.2 torch.cuda API：在代码里“装个监控探头”

光看命令行还不够。很多时候，我们需要在 Python 代码里实时获取显存状态，比如在某个关键步骤前后打点。

PyTorch 提供了一套非常实用的 API，我挑几个最常用的：

import torch

# 当前设备上已分配的显存（字节）
print(torch.cuda.memory_allocated())

# 当前设备上缓存/预留的显存（字节）
print(torch.cuda.memory_reserved())

# 最大峰值显存（从程序开始到现在）
print(torch.cuda.max_memory_allocated())

# 清空缓存（释放未使用的缓存块）
torch.cuda.empty_cache()

这里有个容易混淆的点：memory_allocated 是实际占用的，memory_reserved 是 PyTorch 向 GPU 申请的总量。说白了，reserved 包含了 allocated 加上一些“备用金”。

另外，torch.cuda.empty_cache() 这个函数，我建议只在调试时用。为什么？因为它会强制 PyTorch 释放缓存，但下次需要显存时又得重新申请，反而拖慢速度。生产环境里别乱调。

4.3 内存快照工具：给显存“拍个X光片”

如果前两个工具是“温度计”和“血压计”，那内存快照就是“CT扫描仪”。它能告诉你：每一块显存到底被哪个张量、哪个操作占着。

PyTorch 从 1.8 版本开始，内置了 torch.cuda.memory._snapshot() 功能。用法如下：

# 记录当前时刻的显存快照
snapshot = torch.cuda.memory._snapshot()

# 保存到文件，供后续分析
import pickle
with open('snapshot.pickle', 'wb') as f:
    pickle.dump(snapshot, f)

拿到这个快照文件后，你可以用 PyTorch 官方提供的 内存可视化工具（在 torch.cuda.memory 模块的文档里有链接）来加载查看。它会生成一个火焰图，每个色块代表一个张量或操作，色块越大，占的显存越多。

我记得有一次，一个学员的模型在推理时显存爆了。用 nvidia-smi 看，显存确实满了，但不知道谁干的。后来拍了个快照，发现是 torch.no_grad() 没加对地方，导致中间激活值全部保留了。去掉之后，显存直接降了 40%。你看，这就是快照的威力。

4.4 本章知识体系总览

为了让你更直观地理解这三个工具的关系，我画了一张图：

这三个工具，我建议你按顺序使用：

1. 先跑 nvidia-smi -l 2，看看显存是不是真的满了，哪个进程在吃。
2. 然后在代码里插入 torch.cuda.memory_allocated()，定位到具体是哪一段代码导致的暴涨。
3. 如果还找不到原因，上内存快照，看看是哪个张量在作怪。

说白了，这套组合拳打下来，90% 的显存溢出问题都能现出原形。剩下的 10%，可能是驱动 bug 或者硬件问题，那就得换卡或者升级驱动了。

好了，工具已经给你了。下一章，咱们就拿着这些工具，去实战中抓几个“显存小偷”看看。

---