1、首先查看哪些GPU空闲,nvidia-smi显示当前GPU使用情况。
nvidia-smi
2、然后指定空闲的GPU运行python程序。
CUDA_VISIBLE_DEVICES=0,2,3 python test.py
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2,3"
PS:周期性输出GPU使用情况
每 10s 显示一次GPU使用情况:
watch -n 10 nvidia-smi
安装gpustat库
pip install gpustat
运行命令
gpustat --w
效果如下:
安装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)