python 自动获取可用GPU的ID

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
展开阅读全文
©️2020 CSDN 皮肤主题: 编程工作室 设计师: CSDN官方博客 返回首页
实付0元
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值