1. Take-Away

筆者使用 PyTorch 編寫了不同加速庫(kù)在 ImageNet 上的使用示例（單機(jī)多卡）。需要的同學(xué)可以當(dāng)作 quickstart 將所需要的部分 copy 到自己的項(xiàng)目中（Github 請(qǐng)點(diǎn)擊下面鏈接）：

nn.DataParallel[1] 簡(jiǎn)單方便的 nn.DataParallel

torch.distributed[2] 使用 torch.distributed 加速并行訓(xùn)練

torch.multiprocessing[3] 使用 torch.multiprocessing 取代啟動(dòng)器

apex[4] 使用 apex 再加速

horovod[5] horovod 的優(yōu)雅實(shí)現(xiàn)

注：分布式 evaluation[6]

這里，筆者記錄了使用 4 塊 Tesla V100-PICE 在 ImageNet 進(jìn)行了運(yùn)行時(shí)間的測(cè)試。

測(cè)試結(jié)果發(fā)現(xiàn) Apex 的加速效果最好，但與 Horovod/Distributed 差別不大，平時(shí)可以直接使用內(nèi)置的 Distributed。Dataparallel 較慢，不推薦使用。

（后續(xù)會(huì)補(bǔ)上 V100/K80 上的測(cè)試結(jié)果，穿插了一些試驗(yàn)所以中斷了。）

簡(jiǎn)要記錄一下不同庫(kù)的分布式訓(xùn)練方式，當(dāng)作代碼的 README（我真是個(gè)小機(jī)靈鬼）～

2. 簡(jiǎn)單方便的 nn.DataParallel

DataParallel 可以幫助我們（使用單進(jìn)程控）將模型和數(shù)據(jù)加載到多個(gè) GPU 中，控制數(shù)據(jù)在 GPU 之間的流動(dòng)，協(xié)同不同 GPU 上的模型進(jìn)行并行訓(xùn)練（細(xì)粒度的方法有 scatter，gather 等等）。

DataParallel 使用起來(lái)非常方便，我們只需要用 DataParallel 包裝模型，再設(shè)置一些參數(shù)即可。

需要定義的參數(shù)包括：

參與訓(xùn)練的 GPU 有哪些，device_ids=gpus；

用于匯總梯度的 GPU 是哪個(gè)，output_device=gpus[0] 。

DataParallel 會(huì)自動(dòng)幫我們將數(shù)據(jù)切分 load 到相應(yīng) GPU，將模型復(fù)制到相應(yīng) GPU，進(jìn)行正向傳播計(jì)算梯度并匯總：

model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])

值得注意的是，模型和數(shù)據(jù)都需要先 load 進(jìn) GPU 中，DataParallel 的 module 才能對(duì)其進(jìn)行處理，否則會(huì)報(bào)錯(cuò)：

# 這里要 model.cuda()
model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      # 這里要 images/target.cuda()
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

匯總一下，DataParallel 并行訓(xùn)練部分主要與如下代碼段有關(guān)：

# main.py
import torch
import torch.distributed as dist

gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))

train_dataset = ...

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)

model = ...
model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])

optimizer = optim.SGD(model.parameters())

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

在使用時(shí)，使用 python 執(zhí)行即可：

python main.py

在 ImageNet 上的完整訓(xùn)練代碼，請(qǐng)點(diǎn)擊Github[7]。

3. 使用 torch.distributed 加速并行訓(xùn)練

在 pytorch 1.0 之后，官方終于對(duì)分布式的常用方法進(jìn)行了封裝，支持 all-reduce，broadcast，send 和 receive 等等。通過(guò) MPI 實(shí)現(xiàn) CPU 通信，通過(guò) NCCL 實(shí)現(xiàn) GPU 通信。官方也曾經(jīng)提到用 DistributedDataParallel 解決 DataParallel 速度慢，GPU 負(fù)載不均衡的問(wèn)題，目前已經(jīng)很成熟了～

與 DataParallel 的單進(jìn)程控制多 GPU 不同，在 distributed 的幫助下，我們只需要編寫一份代碼，torch 就會(huì)自動(dòng)將其分配給 n 個(gè)進(jìn)程，分別在 n 個(gè) GPU 上運(yùn)行。

在 API 層面，pytorch 為我們提供了 torch.distributed.launch 啟動(dòng)器，用于在命令行分布式地執(zhí)行 python 文件。在執(zhí)行過(guò)程中，啟動(dòng)器會(huì)將當(dāng)前進(jìn)程的（其實(shí)就是 GPU的）index 通過(guò)參數(shù)傳遞給 python，我們可以這樣獲得當(dāng)前進(jìn)程的 index：

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)

接著，使用 init_process_group 設(shè)置GPU 之間通信使用的后端和端口：

dist.init_process_group(backend='nccl')

之后，使用 DistributedSampler 對(duì)數(shù)據(jù)集進(jìn)行劃分。如此前我們介紹的那樣，它能幫助我們將每個(gè) batch 劃分成幾個(gè) partition，在當(dāng)前進(jìn)程中只需要獲取和 rank 對(duì)應(yīng)的那個(gè) partition 進(jìn)行訓(xùn)練：

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

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

然后，使用 DistributedDataParallel 包裝模型，它能幫助我們?yōu)椴煌?GPU 上求得的梯度進(jìn)行 all reduce（即匯總不同 GPU 計(jì)算所得的梯度，并同步計(jì)算結(jié)果）。all reduce 后不同 GPU 中模型的梯度均為 all reduce 之前各 GPU 梯度的均值：

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

最后，把數(shù)據(jù)和模型加載到當(dāng)前進(jìn)程使用的 GPU 中，正常進(jìn)行正反向傳播：

torch.cuda.set_device(args.local_rank)

model.cuda()

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

匯總一下，torch.distributed 并行訓(xùn)練部分主要與如下代碼段有關(guān)：

# main.py
import torch
import argparse
import torch.distributed as dist

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

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

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

model = ...
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

optimizer = optim.SGD(model.parameters())

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

在使用時(shí)，調(diào)用 torch.distributed.launch 啟動(dòng)器啟動(dòng)：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

在 ImageNet 上的完整訓(xùn)練代碼，請(qǐng)點(diǎn)擊Github[8]。

4. 使用 torch.multiprocessing 取代啟動(dòng)器

有的同學(xué)可能比較熟悉 torch.multiprocessing，也可以手動(dòng)使用 torch.multiprocessing 進(jìn)行多進(jìn)程控制。繞開(kāi) torch.distributed.launch 自動(dòng)控制開(kāi)啟和退出進(jìn)程的一些小毛病～

使用時(shí)，只需要調(diào)用 torch.multiprocessing.spawn，torch.multiprocessing 就會(huì)幫助我們自動(dòng)創(chuàng)建進(jìn)程。

如下面的代碼所示，spawn 開(kāi)啟了 nprocs=4 個(gè)進(jìn)程，每個(gè)進(jìn)程執(zhí)行 main_worker 并向其中傳入 local_rank（當(dāng)前進(jìn)程 index）和 args（即 4 和 myargs）作為參數(shù)：

import torch.multiprocessing as mp

mp.spawn(main_worker, nprocs=4, args=(4, myargs))

這里，我們直接將原本需要 torch.distributed.launch 管理的執(zhí)行內(nèi)容，封裝進(jìn) main_worker 函數(shù)中，其中 proc 對(duì)應(yīng) local_rank（當(dāng)前進(jìn)程 index），進(jìn)程數(shù) nprocs 對(duì)應(yīng) 4， args 對(duì)應(yīng) myargs：

def main_worker(proc, nprocs, args):

   dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)
   torch.cuda.set_device(args.local_rank)

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

   train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

   model = ...
   model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

   optimizer = optim.SGD(model.parameters())

   for epoch in range(100):
      for batch_idx, (data, target) in enumerate(train_loader):
          images = images.cuda(non_blocking=True)
          target = target.cuda(non_blocking=True)
          ...
          output = model(images)
          loss = criterion(output, target)
          ...
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()

在上面的代碼中值得注意的是，由于沒(méi)有 torch.distributed.launch 讀取的默認(rèn)環(huán)境變量作為配置，我們需要手動(dòng)為 init_process_group 指定參數(shù)：

dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)

匯總一下，添加 multiprocessing 后并行訓(xùn)練部分主要與如下代碼段有關(guān)：

# main.py
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

mp.spawn(main_worker, nprocs=4, args=(4, myargs))

def main_worker(proc, nprocs, args):

   dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)
   torch.cuda.set_device(args.local_rank)

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

   train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

   model = ...
   model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

   optimizer = optim.SGD(model.parameters())

   for epoch in range(100):
      for batch_idx, (data, target) in enumerate(train_loader):
          images = images.cuda(non_blocking=True)
          target = target.cuda(non_blocking=True)
          ...
          output = model(images)
          loss = criterion(output, target)
          ...
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()

在使用時(shí)，直接使用 python 運(yùn)行就可以了：

python main.py

在 ImageNet 上的完整訓(xùn)練代碼，請(qǐng)點(diǎn)擊Github[9]。

5.使用 Apex 再加速

Apex 是 NVIDIA 開(kāi)源的用于混合精度訓(xùn)練和分布式訓(xùn)練庫(kù)。Apex 對(duì)混合精度訓(xùn)練的過(guò)程進(jìn)行了封裝，改兩三行配置就可以進(jìn)行混合精度的訓(xùn)練，從而大幅度降低顯存占用，節(jié)約運(yùn)算時(shí)間。此外，Apex 也提供了對(duì)分布式訓(xùn)練的封裝，針對(duì) NVIDIA 的 NCCL 通信庫(kù)進(jìn)行了優(yōu)化。

在混合精度訓(xùn)練上，Apex 的封裝十分優(yōu)雅。直接使用 amp.initialize 包裝模型和優(yōu)化器，apex 就會(huì)自動(dòng)幫助我們管理模型參數(shù)和優(yōu)化器的精度了，根據(jù)精度需求不同可以傳入其他配置參數(shù)。

from apex import amp

model, optimizer = amp.initialize(model, optimizer)

在分布式訓(xùn)練的封裝上，Apex 在膠水層的改動(dòng)并不大，主要是優(yōu)化了 NCCL 的通信。因此，大部分代碼仍與 torch.distributed 保持一致。使用的時(shí)候只需要將 torch.nn.parallel.DistributedDataParallel 替換為 apex.parallel.DistributedDataParallel 用于包裝模型。在 API 層面，相對(duì)于 torch.distributed ，它可以自動(dòng)管理一些參數(shù)（可以少傳一點(diǎn)）：

from apex.parallel import DistributedDataParallel

model = DistributedDataParallel(model)
# # torch.distributed
# model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
# model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

在正向傳播計(jì)算 loss 時(shí)，Apex 需要使用 amp.scale_loss 包裝，用于根據(jù) loss 值自動(dòng)對(duì)精度進(jìn)行縮放：

with amp.scale_loss(loss, optimizer) as scaled_loss:
   scaled_loss.backward()

匯總一下，Apex 的并行訓(xùn)練部分主要與如下代碼段有關(guān)：

# main.py
import torch
import argparse
import torch.distributed as dist

from apex.parallel import DistributedDataParallel

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

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

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

model = ...
model, optimizer = amp.initialize(model, optimizer)
model = DistributedDataParallel(model, device_ids=[args.local_rank])

optimizer = optim.SGD(model.parameters())

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      optimizer.zero_grad()
      with amp.scale_loss(loss, optimizer) as scaled_loss:
         scaled_loss.backward()
      optimizer.step()

在使用時(shí)，調(diào)用 torch.distributed.launch 啟動(dòng)器啟動(dòng)：

UDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

在 ImageNet 上的完整訓(xùn)練代碼，請(qǐng)點(diǎn)擊Github[10]。

6.Horovod 的優(yōu)雅實(shí)現(xiàn)

Horovod 是 Uber 開(kāi)源的深度學(xué)習(xí)工具，它的發(fā)展吸取了 Facebook "Training ImageNet In 1 Hour" 與百度 "Ring Allreduce" 的優(yōu)點(diǎn)，可以無(wú)痛與 PyTorch/Tensorflow 等深度學(xué)習(xí)框架結(jié)合，實(shí)現(xiàn)并行訓(xùn)練。

在 API 層面，Horovod 和 torch.distributed 十分相似。在 mpirun 的基礎(chǔ)上，Horovod 提供了自己封裝的 horovodrun 作為啟動(dòng)器。

與 torch.distributed.launch 相似，我們只需要編寫一份代碼，horovodrun 啟動(dòng)器就會(huì)自動(dòng)將其分配給 n 個(gè)進(jìn)程，分別在 n 個(gè) GPU 上運(yùn)行。在執(zhí)行過(guò)程中，啟動(dòng)器會(huì)將當(dāng)前進(jìn)程的（其實(shí)就是 GPU的）index 注入 hvd，我們可以這樣獲得當(dāng)前進(jìn)程的 index：

import horovod.torch as hvd

hvd.local_rank()

與 init_process_group 相似，Horovod 使用 init 設(shè)置GPU 之間通信使用的后端和端口:

hvd.init()

接著，使用 DistributedSampler 對(duì)數(shù)據(jù)集進(jìn)行劃分。如此前我們介紹的那樣，它能幫助我們將每個(gè) batch 劃分成幾個(gè) partition，在當(dāng)前進(jìn)程中只需要獲取和 rank 對(duì)應(yīng)的那個(gè) partition 進(jìn)行訓(xùn)練：

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

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

之后，使用 broadcast_parameters 包裝模型參數(shù)，將模型參數(shù)從編號(hào)為 root_rank 的 GPU 復(fù)制到所有其他 GPU 中：

hvd.broadcast_parameters(model.state_dict(), root_rank=0)

然后，使用 DistributedOptimizer 包裝優(yōu)化器。它能幫助我們?yōu)椴煌?GPU 上求得的梯度進(jìn)行 all reduce（即匯總不同 GPU 計(jì)算所得的梯度，并同步計(jì)算結(jié)果）。all reduce 后不同 GPU 中模型的梯度均為 all reduce 之前各 GPU 梯度的均值：

hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters(), compression=hvd.Compression.fp16)

最后，把數(shù)據(jù)加載到當(dāng)前 GPU 中。在編寫代碼時(shí)，我們只需要關(guān)注正常進(jìn)行正向傳播和反向傳播：

torch.cuda.set_device(args.local_rank)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

匯總一下，Horovod 的并行訓(xùn)練部分主要與如下代碼段有關(guān)：

# main.py
import torch
import horovod.torch as hvd

hvd.init()
torch.cuda.set_device(hvd.local_rank())

train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

model = ...
model.cuda()

optimizer = optim.SGD(model.parameters())

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
       images = images.cuda(non_blocking=True)
       target = target.cuda(non_blocking=True)
       ...
       output = model(images)
       loss = criterion(output, target)
       ...
       optimizer.zero_grad()
       loss.backward()
       optimizer.step()

在使用時(shí)，調(diào)用 horovodrun 啟動(dòng)器啟動(dòng)：

CUDA_VISIBLE_DEVICES=0,1,2,3 horovodrun -np 4 -H localhost:4 --verbose python main.py

在 ImageNet 上的完整訓(xùn)練代碼，請(qǐng)點(diǎn)擊Github[11]。

7.分布式 evaluation

all_reduce, barrier 等 API 是 distributed 中更為基礎(chǔ)和底層的 API。這些 API 可以幫助我們控制進(jìn)程之間的交互，控制 GPU 數(shù)據(jù)的傳輸。在自定義 GPU 協(xié)作邏輯，匯總 GPU 間少量的統(tǒng)計(jì)信息時(shí)，大有用處。熟練掌握這些 API 也可以幫助我們自己設(shè)計(jì)、優(yōu)化分布式訓(xùn)練、測(cè)試流程。

最近，不少同學(xué)私信了我這樣的問(wèn)題，

訓(xùn)練樣本被切分成了若干個(gè)部分，被若干個(gè)進(jìn)程分別控制運(yùn)行在若干個(gè) GPU 上，如何在進(jìn)程間進(jìn)行通信匯總這些（GPU 上的）信息？

使用一張卡進(jìn)行推理、測(cè)試太慢了，如何使用 Distributed 進(jìn)行分布式地推理和測(cè)試，并將結(jié)果匯總在一起？

......

要解決這些問(wèn)題，我們需要一個(gè)更為基礎(chǔ)的 API，匯總記錄不同 GPU 上生成的準(zhǔn)確率、損失函數(shù)等指標(biāo)信息。這個(gè) API 就是 torch.distributed.all_reduce。

圖2：all_reduce 示意圖

如上圖所示，它的工作過(guò)程包含以下三步：

在調(diào)用 all_reduce(tensor, op=...)后，當(dāng)前進(jìn)程會(huì)向其他進(jìn)程發(fā)送 tensor（例如 rank 0 會(huì)發(fā)送 rank 0 的 tensor 到 rank 1、2、3）

同時(shí)，當(dāng)前進(jìn)程接受其他進(jìn)程發(fā)來(lái)的 tensor（例如 rank 0 會(huì)接收 rank 1 的 tensor、rank 2 的 tensor、rank 3 的 tensor）。

在全部接收完成后，當(dāng)前進(jìn)程（例如rank 0）會(huì)對(duì)當(dāng)前進(jìn)程的和接收到的 tensor （例如 rank 0 的 tensor、rank 1 的 tensor、rank 2 的 tensor、rank 3 的 tensor）進(jìn)行 op （例如求和）操作。

使用 torch.distributed.all_reduce(loss, op=torch.distributed.reduce_op.SUM)，我們就能夠?qū)Σ煌瑪?shù)據(jù)切片（不同 GPU 上的訓(xùn)練數(shù)據(jù)）的損失函數(shù)進(jìn)行求和了。接著，我們只要再將其除以進(jìn)程（GPU）數(shù)量 world_size就可以得到損失函數(shù)的平均值。正確率也能夠通過(guò)同樣方法進(jìn)行計(jì)算：

# 原始代碼
output = model(images)
loss = criterion(output, target)
        
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), images.size(0))
top1.update(acc1.item(), images.size(0))
top5.update(acc5.item(), images.size(0))

# 修改后，同步各 GPU 中數(shù)據(jù)切片的統(tǒng)計(jì)信息，用于分布式的 evaluation
def reduce_tensor(tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.reduce_op.SUM)
    rt /= args.world_size
    return rt

output = model(images)
loss = criterion(output, target)
acc1, acc5 = accuracy(output, target, topk=(1, 5))

torch.distributed.barrier()

reduced_loss = reduce_tensor(loss.data)
reduced_acc1 = reduce_tensor(acc1)
reduced_acc5 = reduce_tensor(acc5)

losses.update(loss.item(), images.size(0))
top1.update(acc1.item(), images.size(0))
top5.update(acc5.item(), images.size(0))

值得注意的是，為了同步各進(jìn)程的計(jì)算進(jìn)度，我們?cè)?reduce 之前插入了一個(gè)同步 API torch.distributed.barrier()。在所有進(jìn)程運(yùn)行到這一步之前，先完成此前代碼的進(jìn)程會(huì)等待其他進(jìn)程。這使得我們能夠得到準(zhǔn)確、有序的輸出。在 Horovod 中，我們無(wú)法使用 torch.distributed.barrier()，取而代之的是，我們可以在 allreduce 過(guò)程中指明：

def reduce_mean(tensor, world_size):
    rt = tensor.clone()
    hvd.allreduce(rt, name='barrier')
    rt /= world_size
    return rt
    
output = model(images)
loss = criterion(output, target)
acc1, acc5 = accuracy(output, target, topk=(1, 5))

reduced_loss = reduce_tensor(loss.data)
reduced_acc1 = reduce_tensor(acc1)
reduced_acc5 = reduce_tensor(acc5)

losses.update(loss.item(), images.size(0))
top1.update(acc1.item(), images.size(0))
top5.update(acc5.item(), images.size(0))

在 ImageNet 上的完整訓(xùn)練代碼，請(qǐng)點(diǎn)擊Github[12]。

8.尾注：

本文中使用的 V100-PICE （前 4 個(gè) GPU）的配置：

圖 3：配置詳情

本文中使用的 V100 （前 4 個(gè) GPU）的配置：

圖 4：配置詳情

本文中使用的 K80 （前 4 個(gè) GPU）的配置：

圖 5：配置詳情

筆者本身是 CV 研究生，今天摸魚的時(shí)候一時(shí)興起研究了一下，后面再慢慢完善～