from gpuinfo import GPUInfo
def get_gpu_id(num_gpu=1):
"""get ID of GPUS
:param num_gpu: num of GPUs to use
:return: gpu_id: ID of allocated GPUs
"""
available_device=GPUInfo.check_empty()
if len(available_device)>=num_gpu:
gpu_id = available_device[:num_gpu]
else:
raise Exception('Only {} GPUs to use!'.format(len(available_device)))
if num_gpu == 1:
gpu_id = gpu_id[0]
return gpu_id
python 自动获取可用GPU的ID
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.