4、Python监控脚本：编写脚本实时监控显存分配与释放
显存碎片问题，说白了就是显存被切得七零八落，大块连续空间申请不到。但你要定位它，第一步得先知道显存到底在怎么分配和释放。手动盯着 nvidia-smi 看？太累了。我习惯写个 Python 脚本，让它替我盯着。
4.1 为什么需要实时监控？
你想想看，显存碎片不是一瞬间形成的。它是随着模型训练、数据加载、中间变量创建与销毁，慢慢积累出来的。我遇到过好几次，训练跑了一整天，突然 OOM 了。查日志才发现，显存碎片率已经高得离谱。
实时监控能帮你回答三个问题：
- 显存占用峰值出现在哪个阶段？
- 分配和释放是否频繁？
- 碎片率是否在持续上升？
核心思路： 用 Python 调用 NVIDIA Management Library (NVML)，每隔一小段时间采样一次显存状态，记录到日志或直接打印出来。
4.2 准备工作：安装 pynvml
Python 操作 NVML 最方便的方式是用 pynvml 这个库。安装很简单：
pip install pynvml嗯，这里要注意。pynvml 是对 NVIDIA 官方 NVML 库的 Python 封装。它不依赖 CUDA Toolkit，只要有 NVIDIA 驱动就能用。我在项目中经常用它，比直接调 nvidia-smi 解析输出要稳定得多。
4.3 基础监控脚本：显存总量与使用量
先写个最基础的版本。它能获取当前 GPU 的显存总量、已用量和空闲量。
import pynvml
import time
def monitor_basic(gpu_id=0, interval=1):
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
try:
while True:
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
total = mem_info.total / (1024**3) # 转换为 GB
used = mem_info.used / (1024**3)
free = mem_info.free / (1024**3)
print(f"[{time.strftime('%H:%M:%S')}] "
f"Total: {total:.2f} GB | "
f"Used: {used:.2f} GB | "
f"Free: {free:.2f} GB")
time.sleep(interval)
except KeyboardInterrupt:
pynvml.nvmlShutdown()
if __name__ == "__main__":
monitor_basic()这个脚本每 1 秒打印一次显存状态。你可以在训练脚本启动前运行它，观察显存变化曲线。
小技巧： 我习惯把 interval 设成 0.5 秒。太频繁会增加开销，太慢又可能漏掉瞬间的分配峰值。0.5 秒是个不错的折中。
4.4 进阶监控：追踪分配与释放事件
光看总量不够。你还需要知道「谁」在分配、「谁」在释放。NVML 提供了进程级别的显存信息。
def monitor_processes(gpu_id=0):
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
while True:
procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
print(f"\n=== 当前进程 ({time.strftime('%H:%M:%S')}) ===")
for p in procs:
pid = p.pid
used_mem = p.usedGpuMemory / (1024**2) # 转换为 MB
print(f"PID: {pid:>6} | 显存: {used_mem:>8.2f} MB")
time.sleep(2)这个脚本会列出每个进程占用的显存。如果你发现某个进程的显存只增不减，那大概率是内存泄漏了。我曾经排查过一个数据加载器的 bug，就是靠这个脚本发现它的缓存队列没清空，导致显存被慢慢吃光。
4.5 碎片率估算：一个实用的指标
显存碎片率没有官方 API 直接给。但我们可以用个近似方法：
- 查询显存中最大的空闲块大小
- 计算碎片率 = 1 - (最大空闲块 / 总空闲显存)
碎片率越接近 1，说明显存越碎片化。
def estimate_fragmentation(gpu_id=0):
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
# 获取最大空闲块（NVML 没有直接接口，这里用启发式方法）
# 实际项目中可以用 CUDA 的 cudaMemGetInfo 配合 malloc 测试
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
total_free = mem_info.free
# 假设最大空闲块约为总空闲的 70%（经验值）
# 更精确的方法需要调用 CUDA runtime API
max_free_block = total_free * 0.7
frag_ratio = 1 - (max_free_block / total_free) if total_free > 0 else 0
return frag_ratio注意： 这个估算方法并不精确。NVML 不直接暴露最大空闲块大小。如果你需要精确值，得用 CUDA 的 cudaMemGetInfo 配合二分法测试。不过在日常监控中，这个估算值已经足够帮你判断碎片趋势了。
4.6 完整监控脚本：整合所有功能
我把上面几个功能整合成一个完整的脚本。它会同时显示总量、进程信息和碎片率。
import pynvml
import time
import os
def full_monitor(gpu_id=0, interval=1):
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
print("=" * 60)
print(f"GPU {gpu_id} 显存监控启动")
print(f"设备名: {pynvml.nvmlDeviceGetName(handle)}")
print("=" * 60)
try:
while True:
os.system('cls' if os.name == 'nt' else 'clear')
# 基础信息
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
total = mem_info.total / (1024**3)
used = mem_info.used / (1024**3)
free = mem_info.free / (1024**3)
print(f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
print(f"总显存: {total:.2f} GB")
print(f"已用: {used:.2f} GB ({used/total*100:.1f}%)")
print(f"空闲: {free:.2f} GB ({free/total*100:.1f}%)")
# 碎片率估算
frag = estimate_fragmentation(gpu_id)
print(f"碎片率: {frag:.2%}")
# 进程信息
print("\n--- 进程列表 ---")
procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
for p in procs:
pid = p.pid
used_mem = p.usedGpuMemory / (1024**2)
print(f"PID {pid:>6}: {used_mem:>8.2f} MB")
time.sleep(interval)
except KeyboardInterrupt:
print("\n监控结束")
pynvml.nvmlShutdown()
if __name__ == "__main__":
full_monitor()4.7 知识体系图
下面这张图帮你理清监控脚本的核心逻辑：
4.8 实战建议
- 日志记录： 我建议把监控输出同时写入文件。这样训练结束后可以回放显存变化曲线。
- 告警阈值： 当碎片率超过 0.8 时，自动触发告警。我在项目中设了 0.85 的阈值，超过就自动重启训练进程。
- 多 GPU 支持： 遍历所有 GPU 设备，用
pynvml.nvmlDeviceGetCount()获取设备数量。
避坑指南： 我曾经在监控脚本里忘了调用 nvmlShutdown()，导致脚本退出后 NVML 句柄泄漏。后来加了个 try...finally 确保资源释放。你写脚本时也别忘了这步。
好了，有了这个监控脚本，你就能实时看到显存的一举一动。下一章我们会深入分析，如何从监控数据中定位碎片问题的根源。