↑ 點(diǎn)擊藍(lán)字?關(guān)注極市平臺(tái)

作者丨縱橫@知乎

來(lái)源丨h(huán)ttps://zhuanlan.zhihu.com/p/98535650

編輯丨極市平臺(tái)

極市導(dǎo)讀

?

利用PyTorch，作者編寫了不同加速庫(kù)在ImageNet上的單機(jī)多卡使用示例，方便讀者取用。

又到適宜劃水的周五啦，機(jī)器在學(xué)習(xí)，人很無(wú)聊。在打開 b 站 “學(xué)習(xí)” 之前看著那空著一半的顯卡決定寫點(diǎn)什么喂飽它們～因此，從 V100-PICE/V100/K80 中各拿出 4 張卡，試驗(yàn)一下哪種分布式學(xué)習(xí)庫(kù)速度最快！這下終于能把剩下的顯存吃完啦，又是老師的勤奮好學(xué)生啦（我真是個(gè)小機(jī)靈鬼）!

Take-Away

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

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

https://github.com/tczhangzhi/pytorch-distributed/blob/master/dataparallel.py

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

https://github.com/tczhangzhi/pytorch-distributed/blob/master/distributed.py

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

https://github.com/tczhangzhi/pytorch-distributed/blob/master/multiprocessing_distributed.py

4、使用 apex 再加速

https://github.com/tczhangzhi/pytorch-distributed/blob/master/apex_distributed.py

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

https://github.com/tczhangzhi/pytorch-distributed/blob/master/horovod_distributed.py

這里，筆者記錄了使用 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ī)靈鬼）～

簡(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。

使用 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)將其分配給 個(gè)進(jì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。

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

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

使用時(shí)，只需要調(diào)用 torch.multiprocessing.spawn，torch.multiprocessing 就會(huì)幫助我們自動(dòng)創(chuàng)建進(jìn)程。如下面的代碼所示，spawn 開啟了 nprocs=4 個(gè)線程，每個(gè)線程執(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），ngpus_per_node 對(duì)應(yīng) 4， args 對(duì)應(yīng) myargs：

def main_worker(proc, ngpus_per_node, 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()

在上面的代碼中值得注意的是，由于沒有 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, ngpus_per_node, 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。

使用 Apex 再加速

Apex 是 NVIDIA 開源的用于混合精度訓(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。

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

Horovod 是 Uber 開源的深度學(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)將其分配給 個(gè)進(jì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。

尾注

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

圖 2：配置詳情

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

圖 3：配置詳情

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

圖 4：配置詳情

筆者本身是 CV 研究生，今天摸魚的時(shí)候一時(shí)興起研究了一下，后面再慢慢完善～工業(yè)界的同學(xué)應(yīng)該有自己的 best practice，feel free to 提 PR 或者留言～

推薦閱讀

• PyTorch多GPU并行訓(xùn)練方法及問(wèn)題整理
  

• PyTorch并行訓(xùn)練指南：?jiǎn)螜C(jī)多卡并行、混合精度、同步BN訓(xùn)練
  

• pytorch多gpu并行訓(xùn)練
  

添加極市小助手微信（ID : cvmart2），備注：姓名-學(xué)校/公司-研究方向-城市（如：小極-北大-目標(biāo)檢測(cè)-深圳），即可申請(qǐng)加入極市目標(biāo)檢測(cè)/圖像分割/工業(yè)檢測(cè)/人臉/醫(yī)學(xué)影像/3D/SLAM/自動(dòng)駕駛/超分辨率/姿態(tài)估計(jì)/ReID/GAN/圖像增強(qiáng)/OCR/視頻理解等技術(shù)交流群：每月大咖直播分享、真實(shí)項(xiàng)目需求對(duì)接、求職內(nèi)推、算法競(jìng)賽、干貨資訊匯總、與?10000+來(lái)自港科大、北大、清華、中科院、CMU、騰訊、百度等名校名企視覺開發(fā)者互動(dòng)交流~

△長(zhǎng)按添加極市小助手

△長(zhǎng)按關(guān)注極市平臺(tái)，獲取最新CV干貨

覺得有用麻煩給個(gè)在看啦~??