PyTorch使用cuda运算

主要使用torch.device()函数选择运行设备。

def try_gpu(i=0):
"""如果存在，则返回gpu(i)，否则返回cpu()"""
if torch.cuda.device_count() >= i + 1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
torch.cuda.device_count()可以返回可用的cuda设备数量，torch.cuda.is_available()可以返回cuda是否可用。

    所属分类：Python     发表于2022-03-08