自动选取可用gpu
pip install nvidia-ml-py==4.304.04
pip安装pynvml pip install nvidia-ml-py3
如果是Python2,则pip nvidia-ml-py2
import os
import pynvml
pynvml.nvmlInit()
def usegpu(need_gpu_count=1):
nouse=[]
for index in range(pynvml.nvmlDeviceGetCount()):
# 这里的0是GPU id
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
used= meminfo.used/meminfo.total
if used<0.5:
nouse.append(index)
if len(nouse)>=need_gpu_count:
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, nouse[:need_gpu_count]))
return nouse[:need_gpu_count]