2、显存监控：用三把尺子量清你的显存

做深度学习的人，谁没被显存坑过？

我刚开始调模型那会儿，经常跑着跑着就 OOM（Out of Memory），然后一脸懵地重启。后来我学乖了——监控显存，就像开车看油表。不看？那你等着半路抛锚吧。

这一章，我带你掌握三套监控工具：nvidia-smi、torch.cuda、py3nvml。它们各有各的脾气，但组合起来，基本能搞定 90% 的显存排查场景。

2.1 nvidia-smi：最直接的命令行工具

说实话，nvidia-smi 是我用得最多的工具。没有之一。因为它简单、快、不用写代码。

打开终端，敲一行命令：

nvidia-smi

你会看到类似这样的输出（我简化了一下关键字段）：

字段	含义	我关注的点
Memory-Usage	显存使用量 / 总显存	如果超过 90%，我会警惕
Volatile GPU-Util	GPU 计算单元利用率	显存高但利用率低？可能数据加载太慢
Processes	占用 GPU 的进程列表	看有没有残留进程在吃显存

我的小习惯： 我会在终端开一个 watch -n 1 nvidia-smi，每秒刷新一次。训练时盯着它，显存涨得不对劲，我立马就能发现。

不过 nvidia-smi 有个缺点——它只能看全局。你想知道「当前 PyTorch 模型占了多少显存」？它帮不了你。这时候就得请出第二把尺子。

2.2 torch.cuda：PyTorch 自带的显存侦探

PyTorch 其实内置了一套显存监控函数。我个人觉得，这套 API 被很多人忽略了。

常用的几个函数：

import torch

# 当前已分配显存（字节）
allocated = torch.cuda.memory_allocated()
print(f"已分配: {allocated / 1024**2:.2f} MB")

# 当前缓存显存（字节）
cached = torch.cuda.memory_reserved()
print(f"已缓存: {cached / 1024**2:.2f} MB")

# 最大分配显存（历史峰值）
max_allocated = torch.cuda.max_memory_allocated()
print(f"峰值: {max_allocated / 1024**2:.2f} MB")

为什么会有「已分配」和「已缓存」的区别？

嗯，这里要注意。PyTorch 为了加速，会多预留一些显存作为缓存。你释放了张量，显存可能不会立刻还给系统，而是留在缓存池里。所以 memory_allocated 才是你模型真正占用的量。

我曾经踩过的坑： 有一次我训练到一半，显存突然飙升。用 nvidia-smi 看，显存快满了。但用 torch.cuda.memory_allocated() 一看，实际分配才 60%。后来发现是 DataLoader 的 num_workers 开太多，每个 worker 都缓存了一堆数据。把 workers 从 8 降到 4，问题解决。

还有一个实用技巧——在代码里加断点，打印显存快照：

# 打印当前所有张量的显存占用
for obj in gc.get_objects():
    if torch.is_tensor(obj) and obj.is_cuda:
        print(f"{type(obj).__name__}: {obj.numel() * obj.element_size() / 1024**2:.2f} MB")

这招在排查「哪个变量没释放」时特别好用。

2.3 py3nvml：更底层的监控库

如果你需要更细粒度的监控——比如看 GPU 温度、风扇转速、功耗——那 nvidia-smi 和 torch.cuda 都不够。这时候我推荐 py3nvml。

它是 NVIDIA Management Library (NVML) 的 Python 封装。说白了，就是直接跟 GPU 驱动对话。

安装很简单：

pip install py3nvml

基本用法：

import py3nvml.py3nvml as nvml

# 初始化
nvml.nvmlInit()

# 获取 GPU 数量
device_count = nvml.nvmlDeviceGetCount()
print(f"检测到 {device_count} 块 GPU")

# 获取第一块 GPU 的信息
handle = nvml.nvmlDeviceGetHandleByIndex(0)

# 显存信息
mem_info = nvml.nvmlDeviceGetMemoryInfo(handle)
print(f"总显存: {mem_info.total / 1024**2:.0f} MB")
print(f"已用显存: {mem_info.used / 1024**2:.0f} MB")
print(f"空闲显存: {mem_info.free / 1024**2:.0f} MB")

# 温度（摄氏度）
temp = nvml.nvmlDeviceGetTemperature(handle, nvml.NVML_TEMPERATURE_GPU)
print(f"GPU 温度: {temp}°C")

# 功耗（毫瓦）
power = nvml.nvmlDeviceGetPowerUsage(handle)
print(f"当前功耗: {power / 1000:.1f} W")

# 清理
nvml.nvmlShutdown()

  什么时候用 py3nvml？ 我个人习惯在写自动化监控脚本时用它。比如训练时每隔 5 秒记录一次显存和温度，画成曲线图。这样训练结束后，我能回看「显存是在哪个 epoch 突然涨上去的」。

不过要注意，py3nvml 需要 root 权限才能读取某些信息（比如风扇转速）。在 Docker 容器里用的话，记得加 --privileged 参数。

2.4 三把尺子怎么选？

我画了张图，帮你理清思路：

说白了，这三把尺子不是互斥的，而是互补的。我自己的习惯是：

日常训练：开一个 watch nvidia-smi 窗口，随时瞄一眼
调试 OOM：在代码里加 torch.cuda.memory_allocated() 打印，定位到具体行
长期跑实验：写个 py3nvml 脚本，每 10 秒记录一次显存+温度，存成 CSV

一个小技巧： 你可以在 PyTorch 训练循环里加个条件：如果 memory_allocated() 超过某个阈值，自动调小 batch size。我有个项目就是这么干的，再也没 OOM 过。

好了，显存监控这块就聊到这儿。下一章我们聊聊怎么用 torch.cuda.empty_cache() 和 del 手动释放显存——别小看这两个操作，用好了能省不少事。

📚 课程目录

2、显存监控：用三把尺子量清你的显存

2.1 nvidia-smi：最直接的命令行工具

2.2 torch.cuda：PyTorch 自带的显存侦探

2.3 py3nvml：更底层的监控库

2.4 三把尺子怎么选？