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)

来源:诸神黄昏的幸存者

物联沃分享整理
物联沃-IOTWORD物联网 » python运行程序设置指定GPU(查看GPU使用情况)

发表评论