nvidia-ml-py更可靠，因其通过NVML直接读取驱动级硬件显存数据（与nvidia-smi一致），涵盖所有进程占用、缓存及碎片，而torch.cuda.memory_allocated()仅统计PyTorch当前Python引用的活跃张量显存，无法反映真实系统压力。
为什么直接用 nvidia-ml-py 获取显存比 torch.cuda.memory_allocated() 更可靠
PyTorch 自身的 torch.cuda.memory_allocated() 只反映当前 Python 引用持有的显存，不包含缓存、未释放的碎片、驱动层保留内存，也不体现其他进程（如 tensorboard、另一个训练脚本）的占用。而 nvidia-ml-py 是 NVIDIA Management Library (NVML) 的 Python 封装，它读取的是 GPU 驱动实时上报的硬件级显存使用量——和 nvidia-smi 显示一致，这才是监控“真实压力”的依据。
常见错误现象：脚本里 torch.cuda.memory_allocated() 显示只用了 2GB，但 nvidia-smi 显示已用 10GB，导致 OOM 被系统 kill；此时靠 PyTorch 自身 API 根本无法预警。
- 必须先安装
nvidia-ml-py：pip install nvidia-ml-py（注意不是nvidia-ml-py3，后者已废弃且不兼容新驱动） - 该库不依赖 PyTorch，即使没 import torch 也能查显存，适合做独立监控进程
- 需要 CUDA 驱动正常工作（
nvidia-smi命令能执行），但不需要 CUDA Toolkit 或 cuDNN
初始化 NVML 并获取单卡显存使用率的最小可行代码
每次调用前都必须初始化 NVML 上下文，且用完需显式 nvmlShutdown()，否则多次运行会报 NVMLError_LibRmVersionMismatch 或内存泄漏。
from pynvml import *
import time
<p>def get_gpu_memory_usage(device_id=0):
try:
nvmlInit()
handle = nvmlDeviceGetHandleByIndex(device_id)
info = nvmlDeviceGetMemoryInfo(handle)</p><h1>info.used / info.total 返回字节值</h1><pre class="brush:php;toolbar:false;"> return info.used / info.total
except NVMLError as e:
print(f"NVML error: {e}")
return None
finally:
try:
nvmlShutdown()
except:
pass示例：每秒打印一次 GPU 0 使用率
while True: usage = get_gpu_memory_usage(0) if usage is not None: print(f"GPU 0: {usage*100:.1f}%") time.sleep(1)
-
device_id从 0 开始编号，对应nvidia-smi左侧的索引 -
info.used是当前被所有进程占用的显存总量（单位字节），info.total是显卡显存总容量 - 不要在循环外调用
nvmlInit()后长期持有句柄——NVML 不是线程安全的，多线程需各自 init/shutdown
监控多卡并触发告警的实用封装
实际训练常跑在多卡环境，需同时监控全部 GPU，并在某卡超阈值时记录日志或发通知。关键点在于：避免因某张卡查询失败（如设备不存在、权限不足）导致整个监控中断。
立即学习“Python免费学习笔记（深入）”；
from pynvml import *
import logging
<p>def check_all_gpus(threshold=0.9):
try:
nvmlInit()
device_count = nvmlDeviceGetCount()
alerts = []
for i in range(device_count):
try:
handle = nvmlDeviceGetHandleByIndex(i)
info = nvmlDeviceGetMemoryInfo(handle)
usage_pct = info.used / info.total
if usage_pct > threshold:
alerts.append(f"GPU{i}: {usage_pct<em>100:.1f}% > {threshold</em>100:.0f}%")
except NVMLError:
alerts.append(f"GPU{i}: unavailable")
return alerts
finally:
try:
nvmlShutdown()
except:
pass</p><h1>使用示例</h1><p>while True:
alerts = check_all_gpus(threshold=0.85)
if alerts:
logging.warning("GPU memory alert: " + "; ".join(alerts))
time.sleep(5)
- 阈值设为 0.85（即 85%）比较稳妥——留出缓冲空间应对突发分配，避免刚到 90% 就触发误报
- PyTorch 分布式训练中，
torch.cuda.device_count()返回的卡数可能少于物理卡数（如某些卡被CUDA_VISIBLE_DEVICES屏蔽），但nvmlDeviceGetCount()始终返回驱动识别到的全部卡数 - 若脚本以非 root 用户运行，某些旧驱动版本可能拒绝返回显存信息，此时需检查
nvidia-smi -q -d MEMORY是否可执行
和 PyTorch 训练循环集成时的注意事项
直接把 get_gpu_memory_usage() 塞进训练 loop 里容易拖慢训练速度，尤其高频调用时。NVML 查询本身耗时约 0.5–2ms，看似不多，但在每个 batch 都查就会累积成显著开销。
- 建议采样频率控制在每 10–30 个 step 一次，或按时间间隔（如每 30 秒）而非 step 触发
- 不要在
torch.no_grad()或eval()模式下调用——这不是原因，而是提醒你：推理阶段同样可能爆显存，监控不能只在 train 模式开启 - 如果训练脚本使用
torch.distributed，确保只在 rank 0 进程中调用 NVML 监控，否则所有进程重复查卡，无意义且增加驱动负载 - 遇到
NVMLError_DriverNotLoaded错误，说明当前环境没有加载 NVIDIA 驱动（比如在 CPU-only 容器里），应捕获并静默跳过，而不是 crash
最易被忽略的一点：NVML 返回的是瞬时快照，它不告诉你哪块显存被谁占了。要定位具体是哪个模型或 dataloader 导致暴涨，仍得配合 torch.cuda.memory_snapshot() + torch.cuda.memory._dump_snapshot() 做事后分析。
