4、显存监控工具：nvidia-smi、nvtop、pyTorch memory profiler的使用方法

显存这东西，说白了就是深度学习里的「硬通货」。模型能不能跑、batch size能开多大、训练会不会OOM，全看它。我见过太多同学代码写得飞起，结果一跑就崩，连哪里爆了都不知道。所以今天咱们就把这三把刀磨利索——nvidia-smi、nvtop、PyTorch Memory Profiler。

4.1 nvidia-smi：最基础的显存快照

nvidia-smi是NVIDIA官方自带的命令行工具，几乎每台有GPU的机器上都有。它的特点是「快照式」——你敲一下命令，它给你拍一张当前显存使用情况的照片。

我个人习惯在训练脚本里每隔几轮就调用一次，看看显存有没有异常增长。比如这样：

import os
import subprocess

def log_gpu_memory():
    result = subprocess.run(['nvidia-smi', '--query-gpu=memory.used,memory.total',
                             '--format=csv,noheader,nounits'],
                            capture_output=True, text=True)
    used, total = result.stdout.strip().split(', ')
    print(f"显存使用: {used} MB / {total} MB")

但要注意，nvidia-smi只能看到全局显存占用，它分不清哪部分是模型参数、哪部分是中间激活值。说白了，它告诉你「水管里有多少水」，但不告诉你「水是从哪个龙头流出来的」。

4.2 nvtop：实时显存「任务管理器」

如果你觉得nvidia-smi太「静态」，那nvtop就是它的升级版。它像Linux里的htop一样，能实时显示每个进程的GPU使用率、显存占用、温度、风扇转速等。

安装很简单：

# Ubuntu/Debian
sudo apt install nvtop

# 或者从源码编译
git clone https://github.com/Syllo/nvtop.git
mkdir build && cd build
cmake ..
make
sudo make install

启动后你会看到一个彩色界面，按F6可以按显存使用量排序。我记得有一次帮同事排查问题，发现他的两个训练进程互相抢显存，就是靠nvtop一眼看出来的——一个进程占了11GB，另一个只分到1GB，两个都在那卡着不动。

4.3 PyTorch Memory Profiler：显微镜级别的分析

前两个工具都是「黑盒」视角。如果你想精确知道——到底是哪一行代码分配了显存？——那就得上PyTorch自带的memory profiler了。

PyTorch提供了torch.cuda.memory_summary()和torch.cuda.memory_snapshot()两个核心API。前者打印摘要，后者生成详细的分配记录。

来看一个实际例子：

import torch

model = MyModel().cuda()
input_tensor = torch.randn(32, 3, 224, 224).cuda()

# 记录当前显存状态
torch.cuda.reset_peak_memory_stats()
output = model(input_tensor)
loss = output.sum()
loss.backward()

# 打印峰值显存
print(torch.cuda.max_memory_allocated() / 1024**3, "GB")

# 打印详细摘要
print(torch.cuda.memory_summary())

输出里会告诉你：当前分配了多少、缓存了多少、峰值是多少、甚至每个张量的大小。我经常用这个来验证我的「显存估算」是否准确——比如我算出来一个batch需要4.2GB，实际跑出来4.5GB，那就说明有0.3GB的额外开销，得去查查是不是哪里多复制了一份数据。

4.4 三把刀怎么配合用？

我一般的工作流是这样的：

1. 日常监控：用nvtop挂在旁边，实时看显存变化曲线。
2. 异常排查：发现显存异常增长，先用nvidia-smi确认是不是只有我的进程在吃显存。
3. 精确定位：用PyTorch Memory Profiler找到具体是哪段代码在「偷」显存。

举个例子，有一次我训练一个Transformer模型，显存每轮都在涨，但模型参数没变。用nvidia-smi一看，确实只有我的进程。再用memory_summary()一查，发现是torch.nn.utils.clip_grad_norm_里创建了一个临时张量没释放。嗯，这就是典型的「显存泄漏」——不是内存泄漏，而是显存。

4.5 知识体系总览

下面这张图总结了三个工具的核心定位和适用场景：

你看，这三个工具其实是层层递进的关系。nvidia-smi是「有没有问题」，nvtop是「谁有问题」，PyTorch Profiler是「问题出在哪一行」。配合使用，基本能覆盖90%的显存排查场景。

想要更多资料，联系可可