pytorch分布式学习笔记

 多设备同步BN

• 不同步时每块GPU单独计算mean/var，batch size/per gpu 小的时候，syncBN的涨点效果明显。
• syncBN的通讯是很费时的（每层BN都要做），eg. 两块GPU训练一个epoch是5.5s，使用syncBN后是8s，而一个GPU训练一个epoch是9s。（来自B站博主霹雳吧啦Wz）

---

DP与DDP

• DP：单进程多线程，只能单机。 【#TODO：进程和线程的区别】
• DDP：多进程多线程，可以多机。在单机情况下速度也比DP快。
• 多GPU训练的启动方式（DDP）：

python -m torch.distributed.launch --help

训练过程中如果强行终止的程序，在开启下次训练前建议通过nvidia-smi指令看下GPU的显存是否全部释放了，有概率出现假死情况。

• DDP的优势：

在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。之后，各进程用该梯度来独立的更新参数。

而 DataParallel是梯度汇总到gpu0，反向传播更新参数，再广播参数给其他的gpu。

总体而言通过减少数据通讯节省了时间。

---

DDP训练过程中的一些代码

Distributed communication package - torch.distributed — PyTorch 1.9.1 documentation  官方代码教程。

分布式加载数据到dataloader：

torch.utils.data.distributed.DistributedSampler ： 多GPU训练时需要将数据分配到各GPU，包括打乱、补齐（复制）、交叉分配。

torch.utils.data.BatchSampler： 将上述分配好的数据打包成batch

    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "val": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}

    # 实例化训练数据集
    train_data_set = MyDataSet(images_path=train_images_path,
                               images_class=train_images_label,
                               transform=data_transform["train"])

    # 实例化验证数据集
    val_data_set = MyDataSet(images_path=val_images_path,
                             images_class=val_images_label,
                             transform=data_transform["val"])

    # 给每个rank对应的进程分配训练的样本索引
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_data_set)
    val_sampler = torch.utils.data.distributed.DistributedSampler(val_data_set)

    # 将样本索引每batch_size个元素组成一个list
    train_batch_sampler = torch.utils.data.BatchSampler(
        train_sampler, batch_size, drop_last=True)

    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    if rank == 0:
        print('Using {} dataloader workers every process'.format(nw))
    train_loader = torch.utils.data.DataLoader(train_data_set,
                                               batch_sampler=train_batch_sampler,
                                               pin_memory=True,
                                               num_workers=nw,
                                               collate_fn=train_data_set.collate_fn)

    val_loader = torch.utils.data.DataLoader(val_data_set,
                                             batch_size=batch_size,
                                             sampler=val_sampler,
                                             pin_memory=True,
                                             num_workers=nw,
                                             collate_fn=val_data_set.collate_fn)

实例化模型并同步BN：

# 实例化模型
model = resnet34(num_classes=num_classes).to(device)

# 只有训练带有BN结构的网络时使用SyncBatchNorm采用意义
# 使用SyncBatchNorm后训练会更耗时
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

# 转为DDP模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

对多块GPU上数据求平均（例如loss）：

“注意，这里对多个设备上的loss求平均不是为了backward，仅仅是查看做个记录。这里有很多人误认为，在使用多GPU时需要先求平均损失然后在反向传播，其实不是的。应该是每个GPU设备计算出各批次数据的损失后，通过backward方法计算出各参数的损失梯度，然后DDP会自动帮我们在多个GPU设备上对各参数的损失梯度求平均，最后通过optimizer.step()去更新各参数”

dist.all_reduce(value)  #此处的value必须是tensor

if average:
    value /= world_size   # 上述代码默认求和，如果是平均需要除以GPU个数
return value

 类似的还有torch.distributed.broadcast，torch.distributed.gather，torch.distributed.reduce，torch.distributed.scatter。scatter可以将信息从master节点传到所有的其他节点，gather可以将信息从别的节点获取到master节点。all_reduce()第一个参数为需要进行运算的变量，第二个参数op则包含了一些方法，例如求和SUM，此外还有MIN, MAX等。broadcast和all_gather好像有点类似，具体不知。注意每个进程上的tensor list长度都必须相同。

【踩坑】上述传递方式的发送变量都是tensor，一些长度不定的list可能很难转成tensor（要补齐什么的）。目前看到的一种非tensor的GPU间传递方法：先存储，后读取。

torch.distributed.broadcast(tensor, src, group=group, async_op=False) # 将tensor从src(rank)广播到group中

torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=group, async_op=False) # 对tensor进行原地in-pllace的reduce，op是torch.distributed.ReduceOp中的一个，指定了某种确定的element-wise的操作

torch.distributed.reduce(tensor, dst, op=ReduceOp.SUM, group=