AI推理显存管理与峰值优化实战 - 第2节
显存监控：三把利刃在手
显存监控这事儿，说白了就是给GPU做体检。我刚开始做大模型训练时，经常遇到OOM（显存溢出）报错，那时候真是两眼一抹黑。后来我养成了一个习惯——训练前先看看显存状态，训练中时不时瞄一眼。嗯，这习惯救了我很多次。
今天我就跟你聊聊三种最常用的监控手段：nvidia-smi、torch.cuda 和 py3nvml。这三把刀各有各的用法，咱们一个一个说。
1. nvidia-smi：命令行里的老司机
nvidia-smi 是NVIDIA自带的工具，简单粗暴。你打开终端，敲一行命令，显存信息就出来了。
nvidia-smi输出大概长这样：
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.86 Driver Version: 470.86 CUDA Version: 11.4 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 Tesla V100-SXM2... Off | 00000000:00:04.0 Off | 0 |
| N/A 34°C P0 37W / 300W| 3245MiB / 32510MiB | 0% Default |
+-------------------------------+----------------------+----------------------+这里有几个关键字段你得注意：
- Memory-Usage：当前用了多少显存，总共多少。比如
3245MiB / 32510MiB，就是用了3.2GB，总共32.5GB。 - GPU-Util：GPU利用率。如果一直0%，说明你的代码没跑在GPU上，或者数据加载太慢。
- Temp：温度。超过80°C就要小心了，我见过有人跑训练跑到90°C，显卡直接降频，速度慢得像蜗牛。
nvidia-smi --query-gpu=memory.used,memory.total --format=csv，输出更干净，适合脚本解析。
1.1 实时监控：watch命令
手动敲命令太累？用 watch 让它自动刷新：
watch -n 1 nvidia-smi这个命令每1秒刷新一次。我一般设成2秒，刷新太快反而看不清变化。
watch 会占用一个终端窗口。如果你在远程服务器上跑实验，建议用 tmux 或 screen 开一个独立窗口来监控。
2. torch.cuda：Python里的内窥镜
如果你在用PyTorch，那 torch.cuda 就是你的贴身小助手。它可以直接在代码里获取显存信息，不用切到终端。
2.1 常用函数一览
| 函数 | 作用 |
|---|---|
torch.cuda.memory_allocated() |
当前已分配的显存（字节） |
torch.cuda.memory_reserved() |
当前缓存池占用的显存（字节） |
torch.cuda.max_memory_allocated() |
从程序开始到现在，分配过的最大显存 |
torch.cuda.memory_summary() |
打印详细的显存使用报告 |
举个例子：
import torch
# 看看当前用了多少显存
allocated = torch.cuda.memory_allocated() / 1024**3 # 转成GB
reserved = torch.cuda.memory_reserved() / 1024**3
print(f"已分配: {allocated:.2f} GB")
print(f"已缓存: {reserved:.2f} GB")
# 打印详细报告
print(torch.cuda.memory_summary())memory_allocated 是实际张量占用的显存，memory_reserved 是PyTorch向GPU申请的缓存池大小。缓存池是为了避免频繁申请释放，但有时候会多占一些显存。
2.2 实战：定位显存泄漏
我曾经遇到过一个坑：训练到第10个epoch时，显存突然爆了。排查了半天，发现是某个循环里创建了张量但没有释放。用 torch.cuda.memory_allocated() 在每个epoch前后打印一下，问题立刻现形。
for epoch in range(100):
before = torch.cuda.memory_allocated()
# 训练代码...
after = torch.cuda.memory_allocated()
print(f"Epoch {epoch}: 显存变化 {after - before:.2f} MB")
if after - before > 100 * 1024**2: # 超过100MB
print("⚠️ 可能有泄漏，检查一下代码！")3. py3nvml：底层监控的瑞士军刀
py3nvml 是NVIDIA Management Library (NVML) 的Python封装。它能获取 nvidia-smi 能拿到的所有信息，甚至更多。比如GPU温度、风扇转速、电源功耗等等。
3.1 安装与基本使用
pip install py3nvml然后写个脚本：
import py3nvml.py3nvml as nvml
# 初始化
nvml.nvmlInit()
# 获取GPU数量
device_count = nvml.nvmlDeviceGetCount()
print(f"检测到 {device_count} 块GPU")
# 遍历每块GPU
for i in range(device_count):
handle = nvml.nvmlDeviceGetHandleByIndex(i)
name = nvml.nvmlDeviceGetName(handle)
mem_info = nvml.nvmlDeviceGetMemoryInfo(handle)
print(f"GPU {i}: {name}")
print(f" 总显存: {mem_info.total / 1024**3:.2f} GB")
print(f" 已用显存: {mem_info.used / 1024**3:.2f} GB")
print(f" 空闲显存: {mem_info.free / 1024**3:.2f} GB")
# 清理
nvml.nvmlShutdown()3.2 高级用法：监控温度与功耗
有一次我在调参时发现训练速度忽快忽慢，用 py3nvml 一查，原来是温度过高导致降频了。
import py3nvml.py3nvml as nvml
nvml.nvmlInit()
handle = nvml.nvmlDeviceGetHandleByIndex(0)
# 获取温度
temp = nvml.nvmlDeviceGetTemperature(handle, nvml.NVML_TEMPERATURE_GPU)
print(f"GPU温度: {temp}°C")
# 获取功耗
power = nvml.nvmlDeviceGetPowerUsage(handle) / 1000 # 转成瓦特
print(f"当前功耗: {power:.2f} W")
# 获取风扇转速
fan_speed = nvml.nvmlDeviceGetFanSpeed(handle)
print(f"风扇转速: {fan_speed}%")
nvml.nvmlShutdown()py3nvml，结果初始化失败。原因是容器没有挂载NVML驱动。解决方案是启动容器时加上 --gpus all 参数，或者挂载 /usr/lib/x86_64-linux-gnu/libnvidia-ml.so。
知识体系总览
下面这张图帮你理清三种监控工具的关系：
总结：怎么选？
这三种工具各有千秋，我的建议是：
- 日常快速查看：用
nvidia-smi，打开终端敲一下就行。 - 代码内调试：用
torch.cuda，尤其是排查显存泄漏时。 - 硬件监控：用
py3nvml，比如你想记录训练过程中的温度变化。
我个人习惯是三个一起用。训练脚本里用 torch.cuda 记录显存变化，同时开一个 watch nvidia-smi 窗口盯着，如果发现异常再用 py3nvml 深挖。嗯，这套组合拳帮我解决了不少显存问题。
记住一句话：显存监控不是目的，目的是让你的模型跑得更稳、更快。别等到OOM了才想起来看显存，那时候就晚了。