Pytorch单机多卡分布式训练 数据并行

Pytorch单机多卡训练（数据并行训练）

Pytorch的数据并行训练，已经被封装的十分完善。全程只需两步：
1.将模型送入多GPU
2.将训练数据送入多GPU

0.判断GPU是否可用

# Device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

1.把模型送入多GPU

if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)
model.to(device)

2.将训练数据送入多GPU

for data in rand_loader:
    input = data.to(device)
    output = model(input)