526互联

pytorch ddp 范例

发布时间 2023-06-09 16:57:44作者: 无左无右

pytorch ddp 范例：

################
## 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.conv2_bn = nn.BatchNorm2d(16, eps=1e-4, momentum=0.01)
        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 = self.pool(F.relu(self.conv2_bn(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)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).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和DDP是没有关系的，所以optimizer初始状态的同一性是不被DDP保证的！
#大多数官方optimizer，其实现能保证从同样状态的model初始化时，其初始状态是相同的。
#所以这边我们只要保证在DDP模型创建后才初始化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原理

  1. 更加深入的了解了下ddp模式下index的分配机制。
    比如总共10个数据, 在程序开始的时候会随机打乱总的indices。
    由于每张卡上打乱的随机种子是相同的，因此可以保证每个进程上的数据集是不重复的，并且能取到所有的数据集。
    A机器数据量10，B机器数据量10，batchsize都是2

master机器分配的：
indices=[4, 7, 3, 0, 6]
Slave机器分配的：
indices=[1, 5, 9, 8, 2]

通过代码加的打印信息如下：

  1. 通过实验得知，DDP模式下都是根据当前机器上面的数据集来确定数据量大小的，只是在划分数据index的时候根据卡数来平分，
indices = indices[self.rank:self.total_size:self.num_replicas]

并且每轮迭代都会重新打乱总的indices。
附pytorch相关源码：
anaconda3/envs/pytorch1.7.0_general/lib/python3.7/site-packages/torch/utils/data/distributed.py

DistributedSampler的__iter__函数
def __iter__(self) -> Iterator[T_co]:
        if self.shuffle:
            """
                由于shuffle=True，因此这一步必定是会执行的
                根据self.epoch + self.seed来确定每一个进程的都是一样的
            """
            # deterministically shuffle based on epoch and seed
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            indices = torch.randperm(len(self.dataset), generator=g).tolist()
        else:
            indices = list(range(len(self.dataset)))

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[:self.total_size]
        assert len(indices) == self.total_size

        # subsample
        # 子采样，指定步长为显卡的数量，根据每张卡的不同次序，指定起点
        # 由于每张卡上打乱的进程是相同的，因此可以保证每个进程上的数据集是不重复的，并且能取到所有的数据集
        indices = indices[self.rank:self.total_size:self.num_replicas]
        # self.num_samples在初始化的时候就已经是所有样本除以进程数量以后的
        # 这里确保取得的索引是和样本数量长度相等的，由于是assert断言，因此必然是相等的
        assert len(indices) == self.num_samples

        return iter(indices)

https://blog.csdn.net/yang332233/article/details/129020200?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522168630008716800213026543%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fblog.%2522%257D&request_id=168630008716800213026543&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2blogfirst_rank_ecpm_v1~rank_v31_ecpm-2-129020200-null-null.268v1koosearch&utm_term=ddp&spm=1018.2226.3001.4450

https://blog.csdn.net/yang332233/article/details/129053867?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522168630012016800215051330%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fblog.%2522%257D&request_id=168630012016800215051330&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2blogfirst_rank_ecpm_v1~rank_v31_ecpm-1-129053867-null-null.268v1koosearch&utm_term=ddp&spm=1018.2226.3001.4450