### 获取 GPU 显存使用情况的方法 要通过 Python 获取 GPU 的显存使用情况，可以借助 `nvidia-smi` 工具或者第三方库如 `pynvml` 或者 `GPUtil` 来实现。 #### 使用 pynvml 库 NVIDIA 提供了一个名为 NVML (NVIDIA Management Library) 的 API 接口来监控 NVIDIA 设备的状态。可以通过安装并调用 `pynvml` 实现这一功能[^1]： ```python import pynvml def get_gpu_memory_usage(): pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 假设只有一张卡，默认索引为0 mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) total_mem = mem_info.total / (1024 ** 3) # 转换为GB单位 used_mem = mem_info.used / (1024 ** 3) free_mem = mem_info.free / (1024 ** 3) return { 'total': f"{total_mem:.2f} GB", 'used': f"{used_mem:.2f} GB", 'free': f"{free_mem:.2f} GB" } gpu_memory = get_gpu_memory_usage() print(f"Total Memory: {gpu_memory['total']}, Used Memory: {gpu_memory['used']}, Free Memory: {gpu_memory['free']}") ``` 上述代码展示了如何初始化 NVML 并读取指定 GPU 卡的内存信息[^1]。 #### 使用 GPUtil 库 另一种更简洁的方式是利用 `GPUtil` 这一轻量级工具包。它封装了许多复杂的底层操作，使得开发者能够更加方便地访问 GPU 数据[^2]: ```python from gpustat import GPUStatCollection gpu_stats = GPUStatCollection.new_query() for gpu in gpu_stats: print(f"ID: {gpu.entry['index']} | " f"Name: {gpu.name} | " f"Load: {gpu.utilization}% | " f"Mem Used: {gpu.memory_used}/{gpu.memory_total} MB") ``` 这里需要注意的是，在实际部署环境中可能需要额外配置权限以便程序正常运行[^2]。 #### 总结 无论是采用官方支持的 `pynvml` 方法还是社区贡献的简化版解决方案比如 `GPUtil`, 都能有效完成对当前系统下 GPU 显存状态查询的任务。具体选择取决于项目需求和个人偏好。