一、引言：显存监控的必要性

在深度学习、科学计算或大规模数据处理场景中，GPU显存的合理分配与实时监控是保障程序稳定运行的关键。显存不足会导致程序崩溃、性能下降甚至硬件损坏，而Python作为主流开发语言，提供了多种查询显存的便捷方式。本文将系统梳理Python中查询显存的核心方法，涵盖命令行工具、深度学习框架内置接口及第三方库，帮助开发者根据需求选择最适合的方案。

二、基础方法：命令行工具NVIDIA-SMI

1.1 NVIDIA-SMI简介

NVIDIA-SMI（System Management Interface）是NVIDIA官方提供的GPU监控工具，可实时获取显存使用量、温度、功耗等关键指标。其优势在于无需安装额外依赖，直接通过命令行调用。

1.2 命令行操作示例

nvidia-smi -l 1  # 每1秒刷新一次GPU状态

输出结果中，Used/Total GPU Memory字段明确显示当前显存使用量与总量。例如：

|   0  Tesla V100-SXM2...  On   | 00000000:00:1E.0 Off |                    0 |
| N/A   35C    P0    65W / 300W |   3421MiB / 32510MiB |      0%      Default |

其中3421MiB / 32510MiB表示已用显存3421MiB，总量32510MiB。

1.3 局限性分析

• 非实时性：命令行调用需手动执行，无法嵌入代码实现自动化监控。
• 数据解析复杂：需通过字符串处理提取显存数值，易出错。

三、深度学习框架内置接口

2.1 PyTorch：torch.cuda模块

PyTorch通过torch.cuda子模块提供显存查询功能，支持精确到当前进程的显存使用统计。

2.1.1 核心方法

import torch
# 查询当前GPU显存总量（单位：MiB）
total_memory = torch.cuda.get_device_properties(0).total_memory / (1024**2)
# 查询当前进程显存分配量（单位：MiB）
allocated_memory = torch.cuda.memory_allocated() / (1024**2)
# 查询当前进程缓存显存量（单位：MiB）
cached_memory = torch.cuda.memory_reserved() / (1024**2)
print(f"Total Memory: {total_memory:.2f} MiB")
print(f"Allocated Memory: {allocated_memory:.2f} MiB")
print(f"Cached Memory: {cached_memory:.2f} MiB")

2.1.2 高级功能：显存快照

PyTorch 1.10+支持torch.cuda.memory_summary()生成详细显存报告，包含内存分配、碎片化程度等数据，适用于调试复杂模型。

2.2 TensorFlow：tf.config模块

TensorFlow通过tf.config.experimental提供显存监控接口，支持多GPU环境下的精细控制。

2.2.1 基础查询

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        details = tf.config.experimental.get_device_details(gpu)
        print(f"Device: {gpu.name}")
        print(f"Total Memory: {details['device_total_memory'] / (1024**2):.2f} MiB")

2.2.2 动态显存分配监控

TensorFlow 2.x支持tf.config.experimental.get_memory_info('GPU:0')实时获取显存使用情况，结合tf.debugging模块可设置显存阈值告警。

四、第三方库：pynvml的进阶应用

3.1 pynvml简介

pynvml是NVIDIA官方NVML库的Python封装，提供比NVIDIA-SMI更底层的控制能力，支持多GPU并行监控与历史数据记录。

3.2 安装与初始化

pip install nvidia-ml-py3

from pynvml import *
nvmlInit()

3.3 核心功能实现

3.3.1 单GPU显存查询

handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print(f"Total: {info.total / (1024**2):.2f} MiB")
print(f"Used: {info.used / (1024**2):.2f} MiB")
print(f"Free: {info.free / (1024**2):.2f} MiB")

3.3.2 多GPU批量监控

device_count = nvmlDeviceGetCount()
for i in range(device_count):
    handle = nvmlDeviceGetHandleByIndex(i)
    info = nvmlDeviceGetMemoryInfo(handle)
    print(f"GPU {i}: Used {info.used / (1024**2):.2f} MiB")

3.3.3 实时监控脚本

结合time.sleep()可实现持续监控：

import time
def monitor_gpu(interval=1):
    try:
        while True:
            for i in range(nvmlDeviceGetCount()):
                handle = nvmlDeviceGetHandleByIndex(i)
                info = nvmlDeviceGetMemoryInfo(handle)
                print(f"GPU {i}: {info.used / (1024**2):.2f} MiB")
            time.sleep(interval)
    except KeyboardInterrupt:
        nvmlShutdown()
monitor_gpu()

五、实战建议：显存管理最佳实践

1. 预分配策略：使用torch.cuda.set_per_process_memory_fraction()限制PyTorch最大显存占用，避免OOM错误。
2. 缓存清理：PyTorch中调用torch.cuda.empty_cache()释放未使用的缓存显存。
3. 日志记录：将pynvml监控数据写入CSV，分析显存使用模式。
4. 异常处理：在训练循环中加入显存检查，超出阈值时自动保存检查点并终止程序。

六、总结与展望

Python查询显存的方法涵盖从简单命令行到深度学习框架内置接口，再到第三方库的完整解决方案。开发者应根据场景选择：

• 快速检查：NVIDIA-SMI
• 框架集成：PyTorch/TensorFlow内置方法
• 精细化控制：pynvml

未来，随着GPU硬件的迭代（如Hopper架构），显存监控工具将进一步支持动态分配、异构计算等高级特性，开发者需持续关注框架更新以优化资源利用率。