2、显存监控：用nvidia-smi和PyTorch的memory_stats()盯紧每一兆显存
显存优化这件事，说白了就是「知己知彼」。你连显存怎么被吃掉的都不知道，谈何优化？
我个人习惯，在开始任何训练之前，第一件事就是打开终端跑一下 nvidia-smi。这就像开车前先看油表，是个基本操作。但光看油表不够，你还得知道发动机内部怎么工作的——这时候PyTorch自带的 memory_stats() 就派上用场了。
2.1 nvidia-smi：最直观的显存仪表盘
nvidia-smi 是NVIDIA官方提供的命令行工具。你直接在终端敲一下，就能看到当前所有GPU的状态。
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 38°C P0 56W / 300W| 10245MiB / 32480MiB | 0% Default |
+-------------------------------+----------------------+----------------------+这里最关键的字段是 Memory-Usage。比如 10245MiB / 32480MiB，意思是你已经用了10GB，总共32GB。嗯，这里要注意：这个数字是「已分配」的显存，不是「实际占用」。什么意思？
我在项目中遇到过好几次：明明模型只占4GB，但 nvidia-smi 显示用了8GB。为什么？因为PyTorch的缓存机制。它为了加速后续分配，会把释放的显存留着不还给系统。这其实是个坑，很多人一看显存爆了就慌了，其实不一定。
torch.cuda.empty_cache()。但这只是清空缓存，不会影响已经分配给张量的显存。
2.2 更精细的监控：nvidia-smi的实时模式
光看一次不够，你得盯着它变化。我习惯用 watch 命令：
watch -n 1 nvidia-smi这样每1秒刷新一次。你训练的时候，就能看到显存怎么一步步涨上去的。如果发现某个瞬间显存突然飙升，那八成是中间变量没释放。
另外，nvidia-smi 还支持查询特定GPU：
nvidia-smi -i 0 # 只看第0号GPU
nvidia-smi --query-gpu=memory.used,memory.free --format=csv # 输出成CSV格式CSV格式特别适合写脚本监控。我曾经写过一个脚本，每5秒记录一次显存使用，然后画成曲线图。这样一眼就能看出训练过程中的显存波动。
2.3 PyTorch的memory_stats()：显微镜级别的分析
nvidia-smi 只能看到「总量」，但你看不到具体是哪个张量、哪个操作在吃显存。这时候就需要PyTorch自带的 torch.cuda.memory_stats() 了。
这个函数返回一个字典，里面包含几十个指标。我挑几个最常用的：
| 指标名 | 含义 | 我的用法 |
|---|---|---|
allocated_bytes.all.current |
当前实际分配的张量占用的显存 | 看模型真正吃了多少 |
reserved_bytes.all.current |
PyTorch向系统申请的显存（含缓存） | 和nvidia-smi对比，看缓存占了多少 |
active_bytes.all.current |
正在被使用的显存（未释放的） | 排查内存泄漏 |
num_alloc_retries |
分配失败后重试的次数 | 如果这个数很大，说明显存快爆了 |
用法很简单：
import torch
# 在训练循环中插入
stats = torch.cuda.memory_stats()
print(f"已分配: {stats['allocated_bytes.all.current'] / 1024**3:.2f} GB")
print(f"已预留: {stats['reserved_bytes.all.current'] / 1024**3:.2f} GB")你想想看，有了这个，你就能精确知道：前向传播吃了多少、反向传播又额外吃了多少。我经常在 backward() 前后各打一次日志，看看梯度计算到底多占显存。
2.4 显存快照：可视化你的显存分配
光看数字不够直观？PyTorch 1.8以后提供了 torch.cuda.memory_snapshot()，可以生成一份完整的显存分配记录。配合 torch.cuda.memory_summary() 能直接打印出漂亮的表格。
print(torch.cuda.memory_summary())输出会显示每个张量的地址、大小、分配栈。嗯，这个信息量很大，但有个问题——它太详细了，几百行输出看得眼花。我一般只在怀疑有内存泄漏时才用。
更实用的做法是：用 torch.cuda.set_per_process_memory_fraction(0.8) 限制PyTorch最多用80%的显存。这样即使有缓存，也不会把显存撑爆。
torch.multiprocessing.set_sharing_strategy('file_system')。
2.5 实战：写一个显存监控装饰器
我习惯把监控逻辑封装成一个装饰器，这样任何函数都能自动打印显存变化：
import torch
import functools
def monitor_memory(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
torch.cuda.reset_peak_memory_stats()
before = torch.cuda.memory_allocated()
result = func(*args, **kwargs)
after = torch.cuda.memory_allocated()
peak = torch.cuda.max_memory_allocated()
print(f"[{func.__name__}] 分配: {(after-before)/1024**2:.2f} MB, "
f"峰值: {peak/1024**2:.2f} MB")
return result
return wrapper
@monitor_memory
def train_step(batch):
# 你的训练逻辑
pass这样每次调用 train_step，都会自动打印显存变化。我在调试新模型时，每个关键函数都加上这个装饰器，很快就能定位到哪个环节在吃显存。
2.6 本章知识体系
下面这张图总结了显存监控的核心工具和它们的关系：
说白了，显存监控就两件事：看总量用nvidia-smi，看细节用memory_stats()。前者帮你快速判断「显存够不够」，后者帮你定位「到底谁在吃显存」。两个工具配合使用，大部分显存问题都能一眼看穿。
嗯，最后提醒一句：别光监控不行动。看到显存异常了，要立刻停下来分析。我曾经有一次盯着监控看了半小时，发现显存稳步上涨，但就是懒得停——结果模型跑到一半OOM了，前面的计算全白费。该停就停，该调就调，这才是工程师该有的态度。