玩弄孕妇人妻系列,无套中出丰满人妻无码99,婷婷亚洲五月色综合,日韩Xx×,99在线精品视频,亚洲V无码,亚洲欧在线观看,日日干夜夜撸

向AI轉(zhuǎn)型的程序員都關(guān)注了這個號??????

機(jī)器學(xué)習(xí)AI算法工程?? 公眾號：datayx

DistributedDataParallel（DDP）是一個支持多機(jī)多卡、分布式訓(xùn)練的深度學(xué)習(xí)工程方法。其能達(dá)到略低于卡數(shù)的加速比,是目前最流行的多機(jī)多卡訓(xùn)練方法。在這篇文章里，作者通過幾個實(shí)例，給大家介紹了DDP在實(shí)際生產(chǎn)中的應(yīng)用，如在DDP中引入SyncBN,多機(jī)多卡環(huán)境下的inference加速等。

零. 概覽

想要讓你的PyTorch神經(jīng)網(wǎng)絡(luò)在多卡環(huán)境上跑得又快又好？那你definitely需要這一篇！

No one knows DDP better than I do!
– – magic_frog（手動狗頭）

本文是DDP系列三篇中的第三篇。本系列力求深入淺出，簡單易懂，猴子都能看得懂（誤）。

基本原理與入門：https://zhuanlan.zhihu.com/p/178402798

實(shí)現(xiàn)原理與源代碼解析：https://zhuanlan.zhihu.com/p/187610959

在過去的兩篇文章里，我們已經(jīng)對DDP的理論、代碼進(jìn)行了充分、詳細(xì)的介紹，相信大家都已經(jīng)了然在胸。但是，實(shí)踐也是很重要的。正所謂理論聯(lián)系實(shí)踐，如果只掌握理論而不進(jìn)行實(shí)踐，無疑是紙上談兵。

在這篇文章里，我們通過幾個實(shí)戰(zhàn)例子，來給大家介紹一下DDP在實(shí)際生產(chǎn)中的應(yīng)用。希望能對大家有所幫助！

在DDP中引入SyncBN
DDP下的Gradient Accumulation的進(jìn)一步加速
多機(jī)多卡環(huán)境下的inference加速
保證DDP性能：確保數(shù)據(jù)的一致性
和DDP有關(guān)的小技巧
控制不同進(jìn)程的執(zhí)行順序
避免DDP帶來的冗余輸出

請歡快地開始閱讀吧！

依賴：pytorch(gpu)>=1.5，python>=3.6

一. 在DDP中引入SyncBN

什么是Batch Normalization(BN)? 這里就不多加以介紹。附上BN文章（https://arxiv.org/abs/1502.03167）。接下來，讓我們來深入了解下BN在多級多卡環(huán)境上的完整實(shí)現(xiàn)：SyncBN。

什么是SyncBN？

SyncBN就是Batch Normalization(BN)。其跟一般所說的普通BN的不同在于工程實(shí)現(xiàn)方式：SyncBN能夠完美支持多卡訓(xùn)練，而普通BN在多卡模式下實(shí)際上就是單卡模式。

我們知道，BN中有moving mean和moving variance這兩個buffer，這兩個buffer的更新依賴于當(dāng)前訓(xùn)練輪次的batch數(shù)據(jù)的計(jì)算結(jié)果。但是在普通多卡DP模式下，各個模型只能拿到自己的那部分計(jì)算結(jié)果，所以在DP模式下的普通BN被設(shè)計(jì)為只利用主卡上的計(jì)算結(jié)果來計(jì)算moving mean和moving variance，之后再廣播給其他卡。這樣，實(shí)際上BN的batch size就只是主卡上的batch size那么大。當(dāng)模型很大、batch size很小時，這樣的BN無疑會限制模型的性能。

為了解決這個問題，PyTorch新引入了一個叫SyncBN的結(jié)構(gòu)，利用DDP的分布式計(jì)算接口來實(shí)現(xiàn)真正的多卡BN。

SyncBN的原理

SyncBN的原理很簡單：SyncBN利用分布式通訊接口在各卡間進(jìn)行通訊，從而能利用所有數(shù)據(jù)進(jìn)行BN計(jì)算。為了盡可能地減少跨卡傳輸量，SyncBN做了一個關(guān)鍵的優(yōu)化，即只傳輸各自進(jìn)程的各自的 小batch mean和 小batch variance，而不是所有數(shù)據(jù)。具體流程請見下面：

前向傳播
在各進(jìn)程上計(jì)算各自的 小batch mean和小batch variance
各自的進(jìn)程對各自的 小batch mean和小batch variance進(jìn)行all_gather操作，每個進(jìn)程都得到s的全局量。
注釋：只傳遞mean和variance，而不是整體數(shù)據(jù)，可以大大減少通訊量，提高速度。
每個進(jìn)程分別計(jì)算總體mean和總體variance，得到一樣的結(jié)果
注釋：在數(shù)學(xué)上是可行的，有興趣的同學(xué)可以自己推導(dǎo)一下。
接下來，延續(xù)正常的BN計(jì)算。
注釋：因?yàn)閺那跋騻鞑サ挠?jì)算數(shù)據(jù)中得到的batch mean和batch variance在各卡間保持一致，所以，running_mean和running_variance就能保持一致，不需要顯式地同步了！
后向傳播：和正常的一樣

貼一下關(guān)鍵代碼，有興趣的同學(xué)可以研究下：pytorch源碼（https://github.com/pytorch/pytorch/blob/release/1.5/torch/nn/modules/_functions.py#L5）

SyncBN與DDP的關(guān)系

一句話總結(jié)，當(dāng)前PyTorch SyncBN只在DDP單進(jìn)程單卡模式中支持。SyncBN用到 all_gather這個分布式計(jì)算接口，而使用這個接口需要先初始化DDP環(huán)境。

復(fù)習(xí)一下DDP的偽代碼中的準(zhǔn)備階段中的DDP初始化階段

d. 創(chuàng)建管理器reducer，給每個parameter注冊梯度平均的hook。
i. 注釋：這一步的具體實(shí)現(xiàn)是在C++代碼里面的，即reducer.h文件。
e. （可能）為可能的SyncBN層做準(zhǔn)備

這里有三個點(diǎn)需要注意：

這里的為可能的SyncBN層做準(zhǔn)備，實(shí)際上就是檢測當(dāng)前是否是DDP單進(jìn)程單卡模式，如果不是，會直接停止。
這告訴我們，SyncBN需要在DDP環(huán)境初始化后初始化，但是要在DDP模型前就準(zhǔn)備好。
為什么當(dāng)前PyTorch SyncBN只支持DDP單進(jìn)程單卡模式？
- 從SyncBN原理中我們可以看到，其強(qiáng)依賴了all_gather計(jì)算，而這個分布式接口當(dāng)前是不支持單進(jìn)程多卡或者DP模式的。當(dāng)然，不排除未來也是有可能支持的。

怎么用SyncBN？

怎么樣才能在我們的代碼引入SyncBN呢？很簡單：

# DDP init
dist.init_process_group(backend='nccl')

# 按照原來的方式定義模型，這里的BN都使用普通BN就行了。
model = MyModel()
# 引入SyncBN，這句代碼，會將普通BN替換成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

# 構(gòu)造DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

又是熟悉的模樣，像DDP一樣，一句代碼就解決了問題。這是怎么做到的呢？

convert_sync_batchnorm的原理：

torch.nn.SyncBatchNorm.convert_sync_batchnorm會搜索model里面的每一個module，如果發(fā)現(xiàn)這個module是、或者繼承了torch.nn.modules.batchnorm._BatchNorm類，就把它替換成SyncBN。也就是說，如果你的Normalization層是自己定義的特殊類，沒有繼承過_BatchNorm類，那么convert_sync_batchnorm是不支持的，需要你自己實(shí)現(xiàn)一個新的SyncBN！

下面給一下convert_sync_batchnorm的源碼（https://github.com/pytorch/pytorch/blob/v1.5.0/torch/nn/modules/batchnorm.py#L474），可以看到convert的過程中，新的SyncBN復(fù)制了原來的BN層的所有參數(shù)：

 @classmethod
def convert_sync_batchnorm(cls, module, process_group=None):
        r"""Helper function to convert all :attr:`BatchNorm*D` layers in the model to
        :class:`torch.nn.SyncBatchNorm` layers.
        """
        module_output = module
if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
            module_output = torch.nn.SyncBatchNorm(module.num_features,
                                                   module.eps, module.momentum,
                                                   module.affine,
                                                   module.track_running_stats,
                                                   process_group)
if module.affine:
with torch.no_grad():
                    module_output.weight = module.weight
                    module_output.bias = module.bias
            module_output.running_mean = module.running_mean
            module_output.running_var = module.running_var
            module_output.num_batches_tracked = module.num_batches_tracked
for name, child in module.named_children():
            module_output.add_module(name, cls.convert_sync_batchnorm(child, process_group))
del module
return module_output

二. DDP下的Gradient Accumulation的進(jìn)一步加速

什么是Gradient Accmulation？

Gradient Accumulation，即梯度累加，相信大家都有所了解，是一種增大訓(xùn)練時batch size的技術(shù)，造福了無數(shù)硬件條件窘迫的我等窮人。不了解的同學(xué)請看這個知乎鏈接（https://www.zhihu.com/question/303070254/answer/573037166）。

為什么還能進(jìn)一步加速？

我們仔細(xì)思考一下DDP下的gradient accumulation。

# 單卡模式，即普通情況下的梯度累加
for 每次梯度累加循環(huán)
    optimizer.zero_grad()
for 每個小step
        prediction = model(data)
        loss_fn(prediction, label).backward()  # 積累梯度，不應(yīng)用梯度改變
    optimizer.step()  # 應(yīng)用梯度改變

我們知道，DDP的gradient all_reduce階段發(fā)生在loss_fn(prediction, label).backward()。這意味著，在梯度累加的情況下，假設(shè)一次梯度累加循環(huán)有K個step，每次梯度累加循環(huán)會進(jìn)行K次 all_reduce！但事實(shí)上，每次梯度累加循環(huán)只會有一次 optimizer.step()，即只應(yīng)用一次參數(shù)修改，這意味著在每一次梯度累加循環(huán)中，我們其實(shí)只要進(jìn)行一次gradient all_reduce即可滿足要求，有K-1次 all_reduce被浪費(fèi)了！而每次 all_reduce的時間成本是很高的！

如何加速

解決問題的思路在于，對前K-1次step取消其梯度同步。幸運(yùn)的是，DDP給我們提供了一個暫時取消梯度同步的context函數(shù) no_sync()（源代碼：https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/distributed.py#L548）。在這個context下，DDP不會進(jìn)行梯度同步。

所以，我們可以這樣實(shí)現(xiàn)加速：

model = DDP(model)

for 每次梯度累加循環(huán)
    optimizer.zero_grad()
# 前K-1個step，不進(jìn)行梯度同步，累積梯度。
for K-1個小step:
with model.no_sync():
            prediction = model(data)
            loss_fn(prediction, label).backward()
# 第K個step，進(jìn)行梯度同步
    prediction = model(data)
    loss_fn(prediction, label).backward()
    optimizer.step()

給一個優(yōu)雅寫法（同時兼容單卡、DDP模式哦）：

from contextlib import nullcontext
# 如果你的python版本小于3.7，請注釋掉上面一行，使用下面這個：
# from contextlib import suppress as nullcontext

if local_rank != -1:
    model = DDP(model)

optimizer.zero_grad()
for i, (data, label) in enumerate(dataloader):
# 只在DDP模式下，輪數(shù)不是K整數(shù)倍的時候使用no_sync
    my_context = model.no_sync if local_rank != -1 and i % K != 0 else nullcontext
with my_context():
        prediction = model(data)
        loss_fn(prediction, label).backward()
if i % K == 0:
        optimizer.step()
        optimizer.zero_grad()

是不是很漂亮！

三. 多機(jī)多卡環(huán)境下的inference加速

問題

有一些非?，F(xiàn)實(shí)的需求，相信大家肯定碰到過：

一般，訓(xùn)練中每幾個epoch我們會跑一下inference、測試一下模型性能。在DDP多卡訓(xùn)練環(huán)境下，能不能利用多卡來加速inference速度呢？
我有一堆數(shù)據(jù)要跑一些網(wǎng)絡(luò)推理，拿到inference結(jié)果。DP下多卡加速比太低，能不能利用DDP多卡來加速呢？

解法

這兩個問題實(shí)際是同一個問題。答案肯定是可以的，但是，沒有現(xiàn)成、省力的方法。

測試和訓(xùn)練的不同在于：

測試的時候不需要進(jìn)行梯度反向傳播，inference過程中各進(jìn)程之間不需要通訊。
測試的時候，不同模型的inference結(jié)果、性能指標(biāo)的類型多種多樣，沒有統(tǒng)一的形式。
我們很難定義一個統(tǒng)一的框架，像訓(xùn)練時model=DDP(model)那樣方便地應(yīng)用DDP多卡加速。

解決問題的思路很簡單，就是各個進(jìn)程中各自進(jìn)行單卡的inference，然后把結(jié)果收集到一起。單卡inference很簡單，我們甚至可以直接用DDP包裝前的模型。問題其實(shí)只有兩個：

我們要如何把數(shù)據(jù)split到各個進(jìn)程中
我們要如何把結(jié)果合并到一起

如何把數(shù)據(jù)split到各個進(jìn)程中：新的data sampler

大家肯定還記得，在訓(xùn)練的時候，我們用的 torch.utils.data.distributed.DistributedSampler幫助我們把數(shù)據(jù)不重復(fù)地分到各個進(jìn)程上去。但是，其分的方法是：每段連續(xù)的N個數(shù)據(jù)，拆成一個一個，分給N個進(jìn)程，所以每個進(jìn)程拿到的數(shù)據(jù)不是連續(xù)的。這樣，不利于我們在inference結(jié)束的時候?qū)⒔Y(jié)果合并到一起。

所以，這里我們需要實(shí)現(xiàn)一個新的data sampler。它的功能，是能夠連續(xù)地劃分?jǐn)?shù)據(jù)塊，不重復(fù)地分到各個進(jìn)程上去。直接給代碼：

# 來源：https://github.com/huggingface/transformers/blob/447808c85f0e6d6b0aeeb07214942bf1e578f9d2/src/transformers/trainer_pt_utils.py
class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):
"""
    Distributed Sampler that subsamples indicies sequentially,
    making it easier to collate all results at the end.
    Even though we only use this sampler for eval and predict (no training),
    which means that the model params won't have to be synced (i.e. will not hang
    for synchronization even if varied number of forward passes), we still add extra
    samples to the sampler to make it evenly divisible (like in `DistributedSampler`)
    to make it easy to `gather` or `reduce` resulting tensors at the end of the loop.
    """

def __init__(self, dataset, batch_size, rank=None, num_replicas=None):
if num_replicas is None:
if not torch.distributed.is_available():
raise RuntimeError("Requires distributed package to be available")
            num_replicas = torch.distributed.get_world_size()
if rank is None:
if not torch.distributed.is_available():
raise RuntimeError("Requires distributed package to be available")
            rank = torch.distributed.get_rank()
self.dataset = dataset
self.num_replicas = num_replicas
self.rank = rank
self.batch_size = batch_size
self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_size
self.total_size = self.num_samples * self.num_replicas

def __iter__(self):
        indices = list(range(len(self.dataset)))
# add extra samples to make it evenly divisible
        indices += [indices[-1]] * (self.total_size - len(indices))
# subsample
        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
return iter(indices)

def __len__(self):
return self.num_samples

如何把結(jié)果合并到一起: all_gather

通過torch.distributed提供的分布式接口all_gather，我們可以把各個進(jìn)程的prediction結(jié)果集中到一起。

難點(diǎn)就在這里。因?yàn)槭澜缟洗嬖谥姘俟值纳窠?jīng)網(wǎng)絡(luò)模型，有著千奇百怪的輸出，所以，把數(shù)據(jù)集中到一起不是一件容易的事情。但是，如果你的網(wǎng)絡(luò)輸出在不同的進(jìn)程中有著一樣的大小，那么這個問題就好解多了。下面給一個方法，其要求網(wǎng)絡(luò)的prediction結(jié)果在各個進(jìn)程中的大小是一模一樣的：

# 合并結(jié)果的函數(shù)
# 1. all_gather，將各個進(jìn)程中的同一份數(shù)據(jù)合并到一起。
#   和all_reduce不同的是，all_reduce是平均，而這里是合并。
# 2. 要注意的是，函數(shù)的最后會裁剪掉后面額外長度的部分，這是之前的SequentialDistributedSampler添加的。
# 3. 這個函數(shù)要求，輸入tensor在各個進(jìn)程中的大小是一模一樣的。
def distributed_concat(tensor, num_total_examples):
    output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(output_tensors, tensor)
    concat = torch.cat(output_tensors, dim=0)
# truncate the dummy elements added by SequentialDistributedSampler
return concat[:num_total_examples]

完整的流程

結(jié)合上面的介紹，我們可以得到下面這樣一個完整的流程。

## 構(gòu)造測試集
# 假定我們的數(shù)據(jù)集是這個
transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
my_testset = torchvision.datasets.CIFAR10(root='./data', train=False, 
        download=True, transform=transform)
# 使用我們的新sampler
test_sampler = SequentialDistributedSampler(my_testset, batch_size=16)
testloader = torch.utils.data.DataLoader(my_testset, batch_size=16, sampler=test_sampler)

# DDP和模型初始化，略。
# ......

# 正式訓(xùn)練和evaluation
for epoch in range(total_epoch_size):
# 訓(xùn)練代碼，略
# .......
# 開始測試
with torch.no_grad():
# 1. 得到本進(jìn)程的prediction
        predictions = []
        labels = []
for data, label in testloader:
            data, label = data.to(local_rank), label.to(local_rank)
            predictions.append(model(data))
            labels.append(label)
# 進(jìn)行g(shù)ather
        predictions = distributed_concat(torch.concat(predictions, dim=0), 
len(test_sampler.dataset))
        labels = distributed_concat(torch.concat(labels, dim=0), 
len(test_sampler.dataset))
# 3. 現(xiàn)在我們已經(jīng)拿到所有數(shù)據(jù)的predictioin結(jié)果，進(jìn)行evaluate！
        my_evaluate_func(predictions, labels)

更簡化的解法

如果我們的目的只是得到性能數(shù)字，那么，我們甚至可以直接在各個進(jìn)程中計(jì)算各自的性能數(shù)字，然后再合并到一起。上面給的解法，是為了更通用的情景。一切根據(jù)你的需要來定！
我們可以單向地把predictions、labels集中到 rank=0的進(jìn)程，只在其進(jìn)行evaluation并輸出。PyTorch也提供了相應(yīng)的接口（鏈接：https://pytorch.org/docs/stable/distributed.html，send和recv）。

四. 保證DDP性能：確保數(shù)據(jù)的一致性

性能期望

從原理上講，當(dāng)沒有開啟SyncBN時，（或者更嚴(yán)格地講，沒有BN層；但一般有的話影響也不大），以下兩種方法訓(xùn)練出來的模型應(yīng)該是性能相似的：

進(jìn)程數(shù)為N的DDP訓(xùn)練
accumulation為N、其他配置完全相同的單卡訓(xùn)練

如果我們發(fā)現(xiàn)性能對不上，那么，往往是DDP中的某些設(shè)置出了問題。在DDP系列第二篇中，我們介紹過一個check list，可以根據(jù)它檢查下自己的配置。其中，在造成性能對不齊的原因中，最有可能的是數(shù)據(jù)方面出現(xiàn)了問題。

DDP訓(xùn)練時，數(shù)據(jù)的一致性必須被保證：各個進(jìn)程拿到的數(shù)據(jù)，要像是accumulation為N、其他配置完全相同的單卡訓(xùn)練中同個accumulation循環(huán)中不同iteration拿到的數(shù)據(jù)。想象一下，如果各個進(jìn)程拿到的數(shù)據(jù)是一樣的，或者分布上有任何相似的地方，那么，這就會造成訓(xùn)練數(shù)據(jù)質(zhì)量的下降，最終導(dǎo)致模型性能下降。

容易錯的點(diǎn)：隨機(jī)數(shù)種子

為保證實(shí)驗(yàn)的可復(fù)現(xiàn)性，一般我們會在代碼在開頭聲明一個固定的隨機(jī)數(shù)種子，從而使得同一個配置下的實(shí)驗(yàn)，無論啟動多少次，都會拿到同樣的結(jié)果。

import random
import numpy as np
import torch

def init_seeds(seed=0, cuda_deterministic=True):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
# Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
if cuda_deterministic:  # slower, more reproducible
        cudnn.deterministic = True
        cudnn.benchmark = False
else:  # faster, less reproducible
        cudnn.deterministic = False
        cudnn.benchmark = True


def main():
# 一般都直接用0作為固定的隨機(jī)數(shù)種子。
    init_seeds(0)

但是在DDP訓(xùn)練中，如果還是像以前一樣，使用0作為隨機(jī)數(shù)種子，不做修改，就會造成以下后果：

DDP的N個進(jìn)程都使用同一個隨機(jī)數(shù)種子
在生成數(shù)據(jù)時，如果我們使用了一些隨機(jī)過程的數(shù)據(jù)擴(kuò)充方法，那么，各個進(jìn)程生成的數(shù)據(jù)會帶有一定的同態(tài)性。
比如說，YOLOv5會使用mosaic數(shù)據(jù)增強(qiáng)（從數(shù)據(jù)集中隨機(jī)采樣3張圖像與當(dāng)前的拼在一起，組成一張里面有4張小圖的大圖）。這樣，因?yàn)楦骺ㄊ褂昧讼嗤碾S機(jī)數(shù)種子，你會發(fā)現(xiàn)，各卡生成的圖像中，除了原本的那張小圖，其他三張小圖都是一模一樣的！
同態(tài)性的數(shù)據(jù)，降低了訓(xùn)練數(shù)據(jù)的質(zhì)量，也就降低了訓(xùn)練效率！最終得到的模型性能，很有可能是比原來更低的。

所以，我們需要給不同的進(jìn)程分配不同的、固定的隨機(jī)數(shù)種子：

def main():
    rank = torch.distributed.get_rank()
# 問題完美解決！
    init_seeds(1 + rank)

五. 和DDP有關(guān)的小技巧

控制不同進(jìn)程的執(zhí)行順序

一般情況下，各個進(jìn)程是各自執(zhí)行的，速度有快有慢，只有在gradient all-reduce的時候，快的進(jìn)程才會等一下慢的進(jìn)程，也就是進(jìn)行同步。那么，如果我們需要在其他地方進(jìn)行同步呢？比如說，在加載數(shù)據(jù)前，如果數(shù)據(jù)集不存在，我們要下載數(shù)據(jù)集：

我們只需要在唯一一個進(jìn)程中開啟一次下載
我們需要讓其他進(jìn)程等待其下載完成，再去加載數(shù)據(jù)

怎么解決這個問題呢？torch.distributed提供了一個barrier()的接口，利用它我們可以同步各個DDP中的各個進(jìn)程！當(dāng)使用barrier函數(shù)時，DDP進(jìn)程會在函數(shù)的位置進(jìn)行等待，知道所有的進(jìn)程都跑到了 barrier函數(shù)的位置，它們才會再次向下執(zhí)行。

只在某進(jìn)程執(zhí)行，無須同步：

這是最簡單的，只需要一個簡單的判斷，用不到barrier()

if rank == 0:
    code_only_run_in_rank_0()

簡單的同步:

沒什么好講的，只是一個示范

code_before()
# 在這一步同步
torch.distributed.barrier()
code_after()

在某個進(jìn)程中執(zhí)行A操作，其他進(jìn)程等待其執(zhí)行完成后再執(zhí)行B操作：

也簡單。

if rank == 0:
    do_A()
    torch.distributed.barrier()
else:
    do_B()
    torch.distributed.barrier()

在某個進(jìn)程中優(yōu)先執(zhí)行A操作，其他進(jìn)程等待其執(zhí)行完成后再執(zhí)行A操作：

這個值得深入講一下，因?yàn)檫@個是非常普遍的需求。利用contextlib.contextmanager，我們可以把這個邏輯給優(yōu)雅地包裝起來！

from contextlib import contextmanager

@contextmanager
def torch_distributed_zero_first(rank: int):
"""Decorator to make all processes in distributed training wait for each local_master to do something.
    """
if rank not in [-1, 0]:
        torch.distributed.barrier()
# 這里的用法其實(shí)就是協(xié)程的一種哦。
yield
if rank == 0:
        torch.distributed.barrier()

然后我們就可以這樣騷操作：

with torch_distributed_zero_first(rank):
if not check_if_dataset_exist():
        download_dataset()
    load_dataset()

優(yōu)雅地解決了需求！

避免DDP帶來的冗余輸出

問題：

當(dāng)我們在自己的模型中加入DDP模型時，第一的直觀感受肯定是，終端里的輸出變成了N倍了。這是因?yàn)槲覀儸F(xiàn)在有N個進(jìn)程在同時跑整個程序。這不光是對有潔癖的同學(xué)造成困擾，其實(shí)對所有人都會造成困擾。因?yàn)楦鱾€進(jìn)程的速度并不一樣快，在茫茫的輸出海洋中，我們難以debug、把控實(shí)驗(yàn)狀態(tài)。

解法：

那么，有什么辦法能避免這個現(xiàn)象呢？下面，筆者給一個可行的方法：logging模塊+輸出信息等級控制。即用logging輸出代替所有print輸出，并給不同進(jìn)程設(shè)置不同的輸出等級，只在0號進(jìn)程保留低等級輸出。舉一個例子：

import logging

# 給主要進(jìn)程（rank=0）設(shè)置低輸出等級，給其他進(jìn)程設(shè)置高輸出等級。
logging.basicConfig(level=logging.INFO if rank in [-1, 0] else logging.WARN)
# 普通log，只會打印一次。
logging.info("This is an ordinary log.")
# 危險(xiǎn)的warning、error，無論在哪個進(jìn)程，都會被打印出來，從而方便debug。
logging.error("This is a fatal log!")