2022年 12月 19日

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)
  • 1
  • 2
  • 3

单GPU:

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

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

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

单GPU:

CUDA_VISIBLE_DEVICES=2 python train.py
  • 1

多GPU:

CUDA_VISIBLE_DEVICES=2,3,4 python train.py
  • 1

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

(1)使用nvidia相关命令

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

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

效果如下:
在这里插入图片描述

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

安装gpustat库。

pip install gpustat
  • 1

运行命令

gpustat --w
  • 1

效果如下:
在这里插入图片描述

(3)使用python的pynvml库

安装pynvml库。

pip install pynvml
  • 1

下面为使用示例:

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)
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7