python运行程序设置指定GPU(查看GPU使用情况)

常见的两种使用方式如下：

一、python文件里设置指定GPU

一般情况：

import os 
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  #（保证程序cuda序号与实际cuda序号对应）
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"  #（代表仅使用第0，1号GPU）

单GPU：

os.environ["CUDA_VISIBLE_DEVICES"] = "2"  #代表只使用第3个gpu

而没设置的时候一般会只用第0号GPU。

二、命令行运行设置指定GPU

单GPU:

CUDA_VISIBLE_DEVICES=2 python train.py

多GPU:

CUDA_VISIBLE_DEVICES=2,3,4 python train.py

附三种查看GPU使用情况的命令：

(1)使用nvidia相关命令

nvidia-smi简称NVSMI，提供监控GPU使用情况和更改GPU状态的功能，是一个跨平台工具，它支持Linux以及从Windows系统。

nvidia-smi
watch -n 1 -d nvidia-smi # 每隔一秒刷新一次

效果如下：

(2)使用gpustat库可实时监测

安装gpustat库。

pip install gpustat

运行命令

gpustat --w

效果如下：

(3)使用python的pynvml库

安装pynvml库。

pip install pynvml

下面为使用示例：

import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 指定显卡号
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total/1024**2) #总的显存大小（float）
print(meminfo.used/1024**2)  #已用显存大小（float）
print(meminfo.free/1024**2)  #剩余显存大小（float）