DiskInfo查看GPU显存占用情况的命令行技巧-平芜编程栈

GPU显存占用监控的命令行实践

在深度学习项目中，模型训练动辄消耗数GB甚至数十GB显存。一个常见的场景是：你信心满满地启动训练脚本，几分钟后却收到“CUDA out of memory”的报错——而此时nvidia-smi显示显存使用量还在稳步上升。这种问题几乎每个AI工程师都遇到过。

更令人困惑的是，有时即使程序已经退出，显存仍未释放；或者多个实验并行时，无法判断哪个进程占用了哪块GPU。面对这些情况，仅靠框架内部的日志远远不够，必须借助系统级工具进行可观测性分析。

本文将从实战角度出发，结合 PyTorch 的内存机制与 CUDA 工具链能力，系统梳理如何通过命令行高效监控和管理 GPU 显存资源。重点不在于介绍某个单一命令，而是构建一套完整的排查逻辑体系。

理解PyTorch的GPU内存管理行为

当我们在 PyTorch 中执行x = torch.randn(1000, 1000).to('cuda')时，表面上只是创建了一个张量，背后其实触发了一整套复杂的内存调度流程。

PyTorch 使用CUDA 缓存分配器（Caching Allocator）来管理显存。它的设计目标不是最小化内存占用，而是最大化分配/释放速度。这意味着：

当一个张量被删除（例如局部变量超出作用域），其占用的显存并不会立即归还给驱动；
这部分内存会被保留在缓存池中，供后续相同或相近大小的请求快速复用；
因此，torch.cuda.memory_allocated()反映的是当前实际使用的显存量，而torch.cuda.memory_reserved()才代表从操作系统角度看已被保留的总量。

import torch if torch.cuda.is_available(): print(f"Device: {torch.cuda.get_device_name(0)}") # 初始状态 print(f"Allocated: {torch.cuda.memory_allocated()/1024**2:.1f} MB") print(f"Reserved: {torch.cuda.memory_reserved()/1024**2:.1f} MB") # 创建大张量 x = torch.randn(20000, 20000, device='cuda') print(f"After allocation → Allocated: {torch.cuda.memory_allocated()/1024**2:.1f} MB") print(f" Reserved: {torch.cuda.memory_reserved()/1024**2:.1f} MB") # 删除张量 del x print(f"After deletion → Allocated: {torch.cuda.memory_allocated()/1024**2:.1f} MB") print(f" Reserved: {torch.cuda.memory_reserved()/1024**2:.1f} MB") # 主动清空缓存（谨慎使用） torch.cuda.empty_cache() print(f"After empty_cache → Reserved: {torch.cuda.memory_reserved()/1024**2:.1f} MB")

运行上述代码会发现：删除张量后，memory_reserved并未下降；只有调用empty_cache()后才真正释放回驱动层。

💡 实践建议：在交互式调试环境中（如 Jupyter Notebook），可以适当调用empty_cache()避免累积占用；但在正式训练脚本中应避免频繁调用，否则会导致反复申请/释放带来性能损耗。

如果你需要详细分析内存分配来源，可使用：

print(torch.cuda.memory_summary(device=None, abbreviated=False))

该命令输出包括按生命周期统计的分配次数、碎片率、最大保留量等信息，对定位显存泄漏非常有帮助。

nvidia-smi：不只是看一眼显存

尽管标题提到 “DiskInfo”，但从上下文来看，这显然是类比表达——意指像查看磁盘信息一样便捷地获取设备状态。真正的标准工具是nvidia-smi，它是深入理解 GPU 使用情况的核心入口。

基础监控：快速掌握全局状态

最简单的命令当然是：

nvidia-smi

它输出的内容远不止显存使用率。关键字段解析如下：

字段	含义
`Memory-Usage`	已用显存 / 总显存（MiB）
`GPU-Util`	GPU 计算单元利用率（非显存带宽）
`Temp`	温度（过高可能触发降频）
`Pwr:Usage/Cap`	功耗使用占比（超限可能导致限速）

值得注意的是，这里的“显存”指的是全局内存（Global Memory），即 GDDR 显存，而非 L1/L2 缓存或共享内存。

持续观测：捕捉动态变化

静态快照难以反映训练过程中的波动。我们可以让nvidia-smi持续刷新：

nvidia-smi -l 1

每秒更新一次，适合观察 batch 处理期间的峰值显存消耗。若发现周期性尖峰接近上限，则说明模型存在瞬时内存压力，可通过梯度累积缓解。

进程级追踪：谁在占用资源？

当你登录服务器却发现 GPU 被占满，但自己并未运行任务时，就需要查清“元凶”。以下命令列出所有正在使用 GPU 的进程：

nvidia-smi --query-compute-apps=pid,name,used_memory,gpu_util --format=csv

输出示例：

pid,name,used_memory [MiB],gpu_util [%] 12345,python,8120 MiB,45% 67890,python,7980 MiB,5%

配合ps命令即可进一步定位具体脚本路径：

ps -p 12345 -o pid,ppid,cmd,%mem,%cpu

结构化输出：用于自动化脚本

如果你想把监控集成到 CI/CD 或训练流水线中，推荐使用 JSON 格式输出：

nvidia-smi -q -D -j

其中-q表示详细查询，-D启用时间戳，-j输出为 JSON。你可以用 Python 解析该结果，实现自动告警或资源预检。

例如，在训练脚本开头加入：

import subprocess import json def get_gpu_memory_used(gpu_id=0): result = subprocess.run(['nvidia-smi', '-q', '-j'], capture_output=True) data = json.loads(result.stdout) mem_used = data['gpu'][gpu_id]['fb_memory_usage']['used'] return mem_used if get_gpu_memory_used() > 10000: # 超过10GB则警告 print("Warning: High GPU memory usage detected!")

这样可以在资源紧张时提前终止，避免 OOM 导致任务中断。

容器环境下的特殊考量

如今大多数团队采用 Docker + PyTorch-CUDA 镜像来统一开发环境。以官方镜像为例：

FROM pytorch/pytorch:2.8-cuda12.1-cudnn8-runtime

这类镜像已预装 PyTorch、CUDA Toolkit 和 cuDNN，极大简化了部署流程。然而，在容器中运行nvidia-smi仍需注意以下几点：

1. 宿主机驱动依赖

虽然容器内无需安装 NVIDIA 驱动，但宿主机必须正确安装对应版本的驱动，并启用 NVIDIA Container Toolkit（原 nvidia-docker）。否则会出现：

Failed to initialize NVML: Driver/library version mismatch

解决方案：确保宿主机驱动版本 ≥ 容器内 CUDA 所需版本。可通过以下命令检查：

# 宿主机执行 cat /proc/driver/nvidia/version nvidia-smi --query-gpu=driver_version --format=csv

2. GPU 可见性控制

默认情况下，容器能看到所有 GPU。为了隔离多用户环境，应通过环境变量限制可见设备：

docker run --gpus '"device=0,1"' -it my-pytorch-image # 或者 CUDA_VISIBLE_DEVICES=0 python train.py

这样即使其他卡被占用，也不会干扰当前任务。

3. 监控穿透问题

由于nvidia-smi读取的是全局设备状态，因此在容器内执行该命令仍能查看所有 GPU 的完整信息——这是正常且必要的行为。但如果希望只关注当前任务所用的卡，建议封装一层过滤逻辑：

# 仅显示第0卡状态 nvidia-smi -i 0

典型问题排查指南

问题一：训练初期就爆显存

现象：模型刚完成前向传播即报错CUDA out of memory。

诊断步骤：
1. 检查模型参数总量：sum(p.numel() for p in model.parameters())
2. 查看 batch size 是否过大，尝试减半测试；
3. 使用torch.cuda.memory_summary()观察前向传播前后显存增长；
4. 若发现某一层骤增，可能是中间激活值过大（如全连接层输入维度爆炸）；

优化手段：
- 改用batch_size=1测试是否通过；
- 启用梯度检查点（Gradient Checkpointing）减少激活存储；
- 使用混合精度训练（torch.cuda.amp）降低张量占用；

from torch.cuda.amp import autocast with autocast(): output = model(input) loss = criterion(output, target) loss.backward()

混合精度可使大部分张量以 FP16 存储，显存占用直接减半。

问题二：重启训练时显存仍被占用

现象：上次训练异常中断，再次启动时报显存不足。

原因分析：Python 进程虽已退出，但 GPU 上下文未完全销毁，导致缓存分配器未释放内存。

解决方法：
1. 执行nvidia-smi查找残留进程 PID；
2. 使用kill -9 <PID>强制终止；
3. 如无明确进程但仍无法分配，尝试重置 GPU 上下文：

# 尝试重置设备（需管理员权限） nvidia-smi --gpu-reset -i 0

⚠️ 注意：此操作会影响该 GPU 上所有任务，请确认无其他重要进程运行。

问题三：多卡训练负载不均

现象：两块 GPU 显存使用分别为 20GB 和 8GB，计算效率低下。

常见原因：
- 数据并行时某些卡负责更多样本；
- 模型结构不对称（如自定义 gather/scatter）；
- 单卡执行额外日志或验证逻辑；

排查方式：
1. 分别在每张卡上打印显存使用：

for i in range(torch.cuda.device_count()): print(f"GPU {i}: {torch.cuda.memory_reserved(i)/1024**2:.1f} MB")

检查DataParallel或DistributedDataParallel是否正确配置；
使用torch.utils.benchmark对数据加载器做压力测试；

理想状态下，各卡显存占用差异应小于 10%。

构建高效的资源管理习惯

掌握工具只是第一步，更重要的是形成良好的工程实践。以下是几个值得采纳的习惯：

✅ 训练脚本开头加入资源检测

import torch def check_gpu_resources(min_free_mb=5000): if not torch.cuda.is_available(): raise RuntimeError("CUDA not available") free_mem = torch.cuda.get_device_properties(0).total_memory \ - torch.cuda.memory_reserved(0) if free_mem < min_free_mb * 1024*1024: raise RuntimeError(f"Not enough GPU memory. Required > {min_free_mb}MB") check_gpu_resources()

避免因环境冲突导致无效训练。

✅ 使用 context manager 自动清理

对于临时占用大量显存的操作，可用上下文管理器包裹：

class gpu_cleanup: def __enter__(self): return self def __exit__(self, *args): torch.cuda.empty_cache() with gpu_cleanup(): large_result = heavy_computation()

确保异常退出时也能释放缓存。

✅ 日志记录关键指标

定期将显存使用写入日志文件，便于事后分析：

import logging logging.basicConfig(filename='training.log', level=logging.INFO) def log_memory(step): alloc = torch.cuda.memory_allocated() / 1024**2 resv = torch.cuda.memory_reserved() / 1024**2 logging.info(f"Step {step}: Allocated={alloc:.1f}MB, Reserved={resv:.1f}MB")

配合绘图工具生成趋势图，直观展现内存增长模式。

这种融合框架 API 与系统工具的监控思路，不仅适用于本地开发，更是大规模分布式训练中不可或缺的一环。随着模型规模持续膨胀，对资源的精细化掌控将成为区分普通开发者与高级工程师的关键能力。