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)
更多推荐
已为社区贡献3条内容
所有评论(0)