监控显存：nvidia-smi命令详解，使用gpustat实时监控，Python脚本监控显存变化

显存监控这事儿，说白了就是给GPU装个"仪表盘"。我刚开始做训练的时候，经常跑着跑着就OOM了，然后一脸懵——到底哪一步爆的？后来养成了习惯，训练前先开监控，心里踏实多了。

今天咱们就聊聊三种监控方式：nvidia-smi、gpustat、还有Python脚本。这三种各有各的用处，我一般搭配着用。

一、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 T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   47C    P0    28W /  70W |   3245MiB / 15109MiB |      0%      Default |
|-------------------------------+----------------------+----------------------+

这里有几个关键字段，我重点说一下：

二、gpustat：更友好的实时监控

nvidia-smi虽然好用，但输出信息太多，看着有点累。gpustat是它的"美化版"，显示更简洁，还带颜色。

安装很简单：

pip install gpustat

使用：

gpustat -i 2  # 每2秒刷新一次

输出效果：

tesla-t4.example.com  Sat Feb 10 14:23:45 2024
[0] Tesla T4 | 47°C, 0% | 3245/15109 MB | user1/python 12345

你看，它把进程信息也列出来了。哪个进程占了多少显存，一目了然。我曾经用这个抓到过同事跑实验忘了关，占了我一半显存。

三、Python脚本监控：把显存数据"抓"到代码里

前面两种都是命令行工具，但有时候我们需要在训练过程中自动记录显存变化，比如画个曲线图，或者当显存接近上限时自动保存模型。这时候就得用Python了。

我常用的方法有两种：

方法1：调用nvidia-smi解析输出

import subprocess
import re

def get_gpu_memory():
    result = subprocess.run(
        ['nvidia-smi', '--query-gpu=memory.used,memory.total', '--format=csv,noheader,nounits'],
        capture_output=True, text=True
    )
    output = result.stdout.strip()
    used, total = output.split(', ')
    return int(used), int(total)

used, total = get_gpu_memory()
print(f"显存使用: {used}MB / {total}MB")

这个方法兼容性好，任何机器只要有nvidia-smi就能用。我在多个服务器上部署过，没出过问题。

方法2：使用pynvml库（推荐）

pip install nvidia-ml-py3

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)

used_mb = info.used // 1024 // 1024
total_mb = info.total // 1024 // 1024
print(f"显存使用: {used_mb}MB / {total_mb}MB")

四、实战：训练时自动记录显存变化

我写了一个小脚本，在训练循环里每100步记录一次显存，训练结束后自动画图。这样就能看到显存是不是在"偷偷"增长——如果有内存泄漏，曲线会一直往上走。

import torch
import pynvml
import matplotlib.pyplot as plt

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)

memory_log = []

for step in range(1000):
    # 训练代码...
    if step % 100 == 0:
        info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        memory_log.append(info.used // 1024 // 1024)

plt.plot(memory_log)
plt.xlabel('Step (x100)')
plt.ylabel('Memory (MB)')
plt.savefig('memory_usage.png')

嗯，这个脚本帮我抓出过好几次显存泄漏。有一次是DataLoader的num_workers设置太大，导致子进程占着显存不释放。还有一次是模型里有个变量忘了detach，梯度图越积越大。

五、知识体系总览

下面这张图总结了本章的核心逻辑：

好了，显存监控这块就聊到这儿。记住一句话：不监控，不训练。养成习惯之后，你会发现OOM不再是噩梦，而是可以预测和预防的。