指定GPU运行python程序

文章目录

• 一、命令行运行python程序时
• 二、在python程序中指定GPU
• 三、使用gpustat库可实时监测
• 四、使用python的pynvml库
• 参考文献

一、命令行运行python程序时

1、首先查看哪些GPU空闲，nvidia-smi显示当前GPU使用情况。

nvidia-smi

• GPU：编号，这里是0和1

• Fan：风扇转速，在0到100%之间变动，第一个是29%

• Name：显卡名，这里两块都是GeForce

• Temp：显卡温度，第一个是60摄氏度

• Perf：性能状态，从P0到P12，P0性能最大，P12最小

• Persistence-M：持续模式的状态开关，该模式耗能大，但是启动新GPU应用时比较快，这里是off

• Pwr：能耗

• Bus-Id：GPU总线

• Disp.A：表示GPU的显示是否初始化

• Memory-Usage：显存使用率

• GPU-Util：GPU利用率，第一个是74%，第二个未用

• Compute M.：计算模式

2、然后指定空闲的GPU运行python程序。

CUDA_VISIBLE_DEVICES=0,2,3 python test.py

二、在python程序中指定GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2,3"

PS：周期性输出GPU使用情况

每 10s 显示一次GPU使用情况：

watch -n 10 nvidia-smi

三、使用gpustat库可实时监测

安装gpustat库

pip install gpustat

运行命令

gpustat --w

效果如下：

四、使用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）

参考文献

• https://www.iotword.com/4559.html