pytorch中分布式训练DDP教程（新手快速入门！）

简介

PyTorch是深度学习领域广泛使用的开源深度学习框架之一。随着深度学习模型的不断增大和数据集的不断增长，单机训练往往不能满足我们的需求。为了加速训练过程，我们可以使用分布式训练技术。在PyTorch中，分布式数据并行（Distributed Data Parallel，简称DDP）是一种常见的分布式训练方法，本篇博客将带您快速上手使用DDP进行分布式训练。

1. DDP简介

DDP是PyTorch中用于在多个GPU或多个计算节点上并行训练的分布式训练技术。它通过数据并行的方式，将模型的参数和梯度分布到多个设备或节点上，从而实现高效的训练。DDP在PyTorch 1.0及以后的版本中得到了原生支持。

2. DDP的优势

使用DDP进行分布式训练有以下几个优势：

a. 加速训练：通过数据并行，DDP能够在多个设备或节点上同时处理不同批次的数据，从而加快训练速度。

b. 内存效率：DDP在每个设备上只保存模型的局部副本和相应的梯度，而不是整个模型的副本，这样可以节省内存。

c. 不需要额外的代码：在PyTorch中，使用DDP进行分布式训练几乎不需要修改您的原始模型和训练代码。

3. 使用DDP进行分布式训练

下面，我们将展示如何在PyTorch中使用DDP进行分布式训练。首先，确保您已经安装了PyTorch的最新版本。

实现DDP流程简单概括如下：
1.首先进行DDP初始化：

dist.init_process_group(backend='nccl')

2.准备数据dataloader和sampler，需要在DDP初始化之后进行：

train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)

3.构造model模型：

model = model.to(local_rank)

4.如果需要Load模型，则要在构造DDP模型之前，且只需要在master上加载就行了：

if dist.get_rank() == 0 and ckpt_path is not None:
    model.load_state_dict(torch.load(ckpt_path))

5.构造DDP model 模型：

model = DDP(model, device_ids=[local_rank], output_device=local_rank)

6.要在构造DDP model之后，才能用model初始化optimizer：

optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

7.创建loss 函数：

loss_func = nn.CrossEntropyLoss().to(local_rank)

8.网络训练，设置DDP sampler的epoch，DistributedSampler需要这个来指定shuffle方式，通过维持各个进程之间的相同随机数种子使不同进程能获得同样的shuffle效果。

trainloader.sampler.set_epoch(epoch)

9.保存模型：save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。因为model其实是DDP model，参数是被**model=DDP(model)**包起来的。并且只需要在进程0上保存一次就行了，避免多次保存重复的东西。

    if dist.get_rank() == 0:
        torch.save(model.module.state_dict(), "%d.ckpt" % epoch)

10.终端bash命令行

# DDP: 使用torch.distributed.launch启动DDP模式
# 使用CUDA_VISIBLE_DEVICES，来决定使用哪些GPU
# CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 main.py

完整的代码如下

################
## main.py文件
import argparse
from tqdm import tqdm
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
# 新增：
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

### 1. 基础模块 ### 
# 假设我们的模型是这个，与DDP无关
class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
# 假设我们的数据是这个
def get_dataset():
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True, 
        download=True, transform=transform)
    # DDP：使用DistributedSampler，DDP帮我们把细节都封装起来了。
    #      用，就完事儿！sampler的原理，第二篇中有介绍。
    train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)
    # DDP：需要注意的是，这里的batch_size指的是每个进程下的batch_size。
    #      也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
    trainloader = torch.utils.data.DataLoader(my_trainset, 
        batch_size=16, num_workers=2, sampler=train_sampler)
    return trainloader
    
### 2. 初始化我们的模型、数据、各种配置  ####
# DDP：从外部得到local_rank参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank

# DDP：DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')  # nccl是GPU设备上最快、最推荐的后端

# 准备数据，要在DDP初始化之后进行
trainloader = get_dataset()

# 构造模型
model = ToyModel().to(local_rank)
# DDP: Load模型要在构造DDP模型之前，且只需要在master上加载就行了。
ckpt_path = None
if dist.get_rank() == 0 and ckpt_path is not None:
    model.load_state_dict(torch.load(ckpt_path))
# DDP: 构造DDP model
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

# DDP: 要在构造DDP model之后，才能用model初始化optimizer。
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# 假设我们的loss是这个
loss_func = nn.CrossEntropyLoss().to(local_rank)

### 3. 网络训练  ###
model.train()
iterator = tqdm(range(100))
for epoch in iterator:
    # DDP：设置sampler的epoch，
    # DistributedSampler需要这个来指定shuffle方式，
    # 通过维持各个进程之间的相同随机数种子使不同进程能获得同样的shuffle效果。
    trainloader.sampler.set_epoch(epoch)
    # 后面这部分，则与原来完全一致了。
    for data, label in trainloader:
        data, label = data.to(local_rank), label.to(local_rank)
        optimizer.zero_grad()
        prediction = model(data)
        loss = loss_func(prediction, label)
        loss.backward()
        iterator.desc = "loss = %0.3f" % loss
        optimizer.step()
    # DDP:
    # 1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
    #    因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
    # 2. 只需要在进程0上保存一次就行了，避免多次保存重复的东西。
    if dist.get_rank() == 0:
        torch.save(model.module.state_dict(), "%d.ckpt" % epoch)


################
## Bash终端命令行运行
# DDP: 使用torch.distributed.launch启动DDP模式
# 使用CUDA_VISIBLE_DEVICES，来决定使用哪些GPU
# CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 main.py

结论

通过本篇博客，我们学习了如何在PyTorch中使用DDP进行分布式训练。DDP可以帮助我们充分利用多个GPU或计算节点的计算资源，加速深度学习模型的训练过程。希望这篇教程能够帮助您快速上手使用DDP进行分布式训练，并提高您的训练效率。在实际应用中，还可以根据需要调整代码以适应更复杂的模型和数据设置。祝您在深度学习的道路上取得更多的成就！