---

前言

随着模型训练需要使用多卡（GPUs）的性能，为此本文记录了关于分布式训练DP和DDP方式。详细介绍了从流程讲解到如何在代码中使用。

---

提示：以下是本篇文章正文内容，下面案例可供参考

一、DP（DataParallel）

DP是单进程多线程实现方式。

假设参数量为M，节点数为N.
GUP0:
传入梯度为：（N-1）M
传出参数为：（N-1）M
其他GPU:
传出梯度为M
传入参数为M
缺点：
单进程，多线程，Python GIL只能利用一个CUP核。
GPU0负责收集梯度，更新参数，同步参数。通信，计算压力大。

二、DP使用

# 以yolov5为例
if cuda and RANK == -1 and torch.cuda.device_count()>1:
	LOGGER.warning('WARNING: ......')
	model = torch.nn.DataParallel(model)
# 以自己搭建模型为例
import torch
from torch import nn
from Net import model_

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if torch.cuda.device_count()>1:
	model = torch.nn.DataParallel(model_)

三、DDP（DistributedDataParallel）

DDP是多进程工作方式，针对多机多卡（也包括单机多卡）。与DP不同的是，不需要全程使用主GPU来汇总和分发数据。

传递数据，求和，形成一个闭环。使得每张卡同时再收发数据。
Scatter-Reduce:通过分发让每个节点有不同的值，并计算求和。


AllGather:使得每张卡上有相同的梯度数据。

总体流程：

DDP总结：
各个GPU独立从硬盘加载输入数据，加载相同的模型，数据均匀划分；
各个GPU独立的进行前向传播，损失计算，反向传播；
每两个相邻的GOU进行梯度的汇总和分发；
各个GPU进行独立的更新参数；
核心：
分为多个桶，将前向传播梯度每满一个桶时候，在进行一次Ring-AllReduce;同时每个GPU仍在进行梯度计算。


当所有的桶都同步之后（梯度相同），每个GPU再使用各自的优化器等进行优化反向传播。

四、DDP使用

1.准备DDP环境
在使用DDP训练时，首先要初始化一下DDP环境，设置后通信后端，进程组等。如下代码所示：

def setup(rank, world_size):
	# MASTER_ADDR指定负责协调初始化过程的主机地址，这是表示单机多GPU的所有进程都将连接到本地机器上
	os.environ['MASTER_ADDR'] = 'localhost'
	#MASTER_PORT制定了主机监听的端口号，用于进程间通信。
	os.environ['MASTER_PORT'] = '123456'
	#初始化分布式进程组。rank是当前进程再进程组中的编号，world_size是总进程数（GPU数量）即进程组大小。
	dist.init_process_grop(backend=('nccl' if dist.is_nccl_available() else 'gloo'), rank=rank, world_size=world_size)

2.准备数据加载器

def get_loader(trainset, testset, batch_size, rank, world_size):
    train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank)
	train_loader = DataLoader(train_set, batch_size=batch_size, sampler=train_sampler)
    # 对于测试集来说，可以选择使用DistributedSampler，也可以选择不使用，这里选择使用
    test_sampler = DistributedSampler(test_set, num_replicas=world_size, rank=rank)
	test_loader = DataLoader(test_set, batch_size=batch_size, sampler=train_sampler)
    # 不使用的代码很简单， 如下所示
    # test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)
    rerturn train_loader, test_loader

3.准备DDP模型和优化器
在定义好模型之后，需要再所有进程中复制模型并用DDP封装，代码如下：

def prepare_model_and_optimizer(model, rank, lr):
    # 设置设备为当前进程的GPU。这里`rank`代表当前进程的编号，
    # `cuda:{rank}` 指定模型应该运行在对应编号的GPU上。
    device = torch.device(f"cuda:{rank}") 
    # 包装模型以使用分布式数据并行。DDP将在多个进程间同步模型的参数，
    # 并且只有指定的`device_ids`中的GPU才会被使用。
    model = model.to(device)
    model = DDP(model, device_ids=[rank])
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    return model, optimizer

DDP中不同进程之间指挥同步梯度，为了确保训练时的参数同步，需要再训练开始前确保不同进程上模型和优化器的初始状态相同。确保模型的初始状态相同有如下两种方式：
1）参数初始化方法
每个GPU上都会有一个模型，为此利用统一的初始化方法，确保不同GPU上的参数统一性。

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
    elif isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)
torch.manual_seed(42)  # 设置随机种子以确保可重复性
# 设置所有GPU的随机种子
torch.cuda.manual_seed_all(42)
model = MyModel()
model.apply(weights_init)

2）加载相同的与训练权重

model = MyModel()
model.load_state_dict(torch.load("path_to_weights.pth"))

注：如果没有设置初始化方法同时也没有模型权重。一个可行方法是手动同步，将rank=0的进程上模型文件临时保存，然后其他进程加载，最后删除临时文件。

def synchronize_model(model, rank, root='temp_model.pth'):
    if rank == 0:
        # 保存模型到文件
        torch.save(model.state_dict(), root)
    torch.distributed.barrier()  # 等待rank=0保存模型

    if rank != 0:
        # 加载模型权重
        model.load_state_dict(torch.load(root))
    torch.distributed.barrier()  # 确保所有进程都加载了模型

    if rank == 0:
        # 删除临时文件
        os.remove(root)

模型同步可以省略，再使用torch.nn.parallel.DistributeDataParallel封装模型时，它会在内部处理所需的同步操作。
4.开始训练
训练时的代码，和单卡没有什么区别。最主要的就是再每个epoch开始的时候，要设置一下sampler的epoch，确保每个epoch的残阳数据的顺序是不一样的。

def train(model, optimizer, criterion, rank, train_loader, num_epochs):
    sampler = train_loader.sampler
    for epoch in range(num_epochs):
        # 在每个epoch开始时更新sampler
        sampler.set_epoch(epoch)
        model.train()
        for batch_idx, (data, targets) in enumerate(dataloader):
            data, targets = data.cuda(rank), targets.cuda(rank)
            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            # 只在rank为0的进程中打印信息
            if rank == 0 and batch_idx % 100 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}")

注：这里打印的loss只是rank0上的loss,如果要打印所有卡上的平均loss,需要使用all_reduce方法。

	# 将损失从所有进程中收集起来并求平均
    # 创建一个和loss相同的tensor，用于聚合操作
    reduced_loss = torch.tensor([loss.item()]).cuda(rank)
    # all_reduce操作默认是求和
    dist.all_reduce(reduced_loss)
    # 求平均
    reduced_loss = reduced_loss / dist.get_world_size()

    # 只在rank为0的进程中打印信息
    if rank == 0 and batch_idx % 100 == 0:
        print(f"Epoch {epoch}, Batch {batch_idx}, Average Loss: {reduced_loss.item()}")

评估代码如果使用了DistributedSampler，计算指标时，需要gather每个进程上的preds和gts，然后计算全局指标。如果没有设置DistributedSampler，和单卡代码一致即可。

def evaluate(model, test_loader, rank):
    model.eval()
    total_preds = []
    total_targets = []

    with torch.no_grad():
        for data, targets in test_loader:
            data, targets = data.to(rank), targets.to(rank)
            outputs = model(data)
            _, preds = torch.max(outputs, 1)

            # 收集当前进程的结果
            total_preds.append(preds)
            total_targets.append(targets)

    # 将所有进程的preds和targets转换为全局列表
    total_preds = torch.cat(total_preds).cpu()
    total_targets = torch.cat(total_targets).cpu()

    # 使用all_gather将所有进程的数据集中到一个列表中
    gathered_preds = [torch.zeros_like(total_preds) for _ in range(dist.get_world_size())]
    gathered_targets = [torch.zeros_like(total_targets) for _ in range(dist.get_world_size())]
    
    dist.all_gather(gathered_preds, total_preds)
    dist.all_gather(gathered_targets, total_targets)

    if rank == 0:
        # 只在一个进程中进行计算和输出
        gathered_preds = torch.cat(gathered_preds)
        gathered_targets = torch.cat(gathered_targets)
        
        # 计算全局性能指标
        accuracy = (gathered_preds == gathered_targets).float().mean()
        print(f'Global Accuracy: {accuracy.item()}')

---

总结

参考文章：DDP