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)