pytorch使用多GPU

# 查看几块GPU
import torch
print(torch.cuda.device_count())

# 查看显卡配置信息
nvidia-smi

# 单主机多块GPU使用
device = torch.device("cuda:0" if torch.cuda.is_avaliable() else :"cpu")
device0 = torch.device("cuda:0")
device1 = torch.device("cuda:1")

net = torch.nn.DataParallel(model)  # 可用的GPU都可以被使用
net.to(device)

# 如何利用部分GPU ？？？？？
# 假设有4块GPU 【0,1,2,3】
device_id = [0, 1, 2, 3]
input_data = input_data.to(device=device_id[0])
net = torch.nn.DataParallel(model)
net.to(device)

Pytorch使用GPU进行训练注意事项 | 文艺数学君