pytorch使用记录（六） 明确使用哪块GPU

可以用在多GPU的服务器上，明确使用的是哪块GPU，可以将信息打印出来，返回的是device的列表。

具体用到的pytorch函数是torch.cuda.device_count()，返回可得到的GPU数量。

接下来是具体函数代码，分为两部分，一个主函数get_proper_device明确使用GPU还是CPU，另外一个函数get_proper_cuda_device是明确具体使用哪块GPU。

# 获取 GPU信息
def get_proper_cuda_device(device, verbose=True):
    if not isinstance(device, list):
        device = [device]
    count = torch.cuda.device_count()
    if verbose:
        print("[Builder]: Found {} gpu".format(count)) 
    for i in range(len(device)):
        d = device[i]
        d_id = None
        if isinstance(d, str):
            #  正则表达式，查看是否存在“cuda:0”这种形式。
            if re.search("cuda:[d]+", d):
                d_id = int(d[5:])
        elif isinstance(d, int):
            d_id = d
        if d_id is None:
            raise ValueError("[Builder]: Wrong cuda id {}".format(d))
        if d_id < 0 or d_id >= count:
            if verbose:
                print("[Builder]: {} is not found, ignore.".format(d))
            device[i] = None
        else:
            device[i] = d_id
    device = [d for d in device if d is not None]  # 返回可用的GPU，去掉None元素。
    return device

# 确定使用GPU还是CPU
def get_proper_device(devices, verbose=True):
    devices = copy.copy(devices)  # 浅拷贝，避免改变传入参数的值
    if not isinstance(devices, list):
        devices = [devices]
    use_cpu = any([d.find("cpu")>=0 for d in devices])
    use_gpu = any([(d.find("cuda")>=0 or isinstance(d, int)) for d in devices])
    assert not (use_cpu and use_gpu), "{} contains cpu and cuda device.".format(devices)
    if use_gpu:
        devices = get_proper_cuda_device(devices, verbose)
        if len(devices) == 0:
            if verbose:
                print("[Builder]: Failed to find any valid gpu in {}, use `cpu`.".format(origin))
            devices = ["cpu"]
    return devices

个人觉得这部分代码的作用是对起始的devices列表中的cuda编号进行检查，对出现的问题进行打印报错。返回的device是只有数字编号的列表。

其实如果想要简单使用的话，可以参考pytorch使用记录（三） 多GPU训练。

此外，发现一个对多GPU训练介绍很详细的博客，可以看一下。Pytorch中多GPU训练指北。介绍了Dataparallel和DistributedParallel，并对pin_memory锁页内存进行介绍。

最后

以上就是独特心锁为你收集整理的pytorch使用记录（六） 明确使用哪块GPU的全部内容，希望文章能够帮你解决pytorch使用记录（六） 明确使用哪块GPU所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。