从零搭建Pytorch模型教程（七）单机多卡和多机多卡训练

前言

本文主要介绍单机多卡训练和多机多卡训练的实现方法和一些注意事项。其中单机多卡训练介绍两种实现方式，一种是DP方式，一种是DDP方式。多机多卡训练主要介绍两种实现方式，一种是通过horovod库，一种是DDP方式。

单机单卡训练

前面我们已经介绍了一个完整的训练流程，但这里由于要介绍单机多卡和多机多卡训练的代码，为了能更好地理解它们之间的区别，这里先放一个单机单卡也就是一般情况下的代码流程。

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_dataset = ...
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)

model = ...
optimizer = optim.SGD(model.parameters())

for epoch in range(opt.num_epoch):
   for i, (input, target) in enumerate(train_loader):
      input= input.to(device)
      target = target.to(device)
      ...
      output = model(input)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

单机多卡训练

单机多卡训练的部分有两种实现方式，一种是DP方式，一种是DDP方式。

nn.DataParallel(DP)
DP方式比较简单，仅仅通过nn.DataParallel对网络进行处理即可。

其它部分基本与单机单卡训练流程相同。

import torch

train_dataset = ...
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)

model = ...
model = nn.DataParallel(model.to(device)<