並行訓練方法

Take-Away

筆者使用 PyTorch 編寫了不同加速庫在 ImageNet 上的使用示例（單機多卡），需要的同學可以當作 quickstart 將需要的部分 copy 到自己的項目中（Github 請點擊下面鏈接）：

nn.DataParallel簡單方便的 nn.DataParallel

2. torch.distributed 使用 torch.distributed 加速並行訓練

3. torch.multiprocessing 使用 torch.multiprocessing 取代啓動器

4. apex 使用 apex 再加速

5. horovod horovod 的優雅實現

這裏，筆者記錄了使用 4 塊 Tesla V100-PICE 在 ImageNet 進行了運行時間的測試，測試結果發現 Apex 的加速效果最好，但與 Horovod/Distributed 差別不大，平時可以直接使用內置的 Distributed。Dataparallel 較慢，不推薦使用。（後續會補上 V100/K80 上的測試結果，穿插了一些試驗所以中斷了）

圖 1：在 ImageNet 2012 上訓練和測試一個 epoch 所需的時間（V100-PICE）

簡要記錄一下不同庫的分佈式訓練方式，當作代碼的 README（我真是個小機靈鬼）～

簡單方便的 nn.DataParallel

DataParallel 可以幫助我們（使用單進程控）將模型和數據加載到多個 GPU 中，控制數據在 GPU 之間的流動，協同不同 GPU 上的模型進行並行訓練（細粒度的方法有 scatter，gather 等等）。

DataParallel 使用起來非常方便，我們只需要用 DataParallel 包裝模型，再設置一些參數即可。需要定義的參數包括：參與訓練的 GPU 有哪些，device_ids=gpus；用於彙總梯度的 GPU 是哪個，output_device=gpus[0] 。DataParallel 會自動幫我們將數據切分 load 到相應 GPU，將模型複製到相應 GPU，進行正向傳播計算梯度並彙總：

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

值得注意的是，模型和數據都需要先 load 進 GPU 中，DataParallel 的 module 才能對其進行處理，否則會報錯：

# 這裏要 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 並行訓練部分主要與如下代碼段有關：

# 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()

在使用時，使用 python 執行即可：

python main.py

在 ImageNet 上的完整訓練代碼，請點擊Github。

使用 torch.distributed 加速並行訓練

在 pytorch 1.0 之後，官方終於對分佈式的常用方法進行了封裝，支持 all-reduce，broadcast，send 和 receive 等等。通過 MPI 實現 CPU 通信，通過 NCCL 實現 GPU 通信。官方也曾經提到用 DistributedDataParallel 解決 DataParallel 速度慢，GPU 負載不均衡的問題，目前已經很成熟了～

與 DataParallel 的單進程控制多 GPU 不同，在 distributed 的幫助下，我們只需要編寫一份代碼，torch 就會自動將其分配給個進程，分別在個 GPU 上運行。

在 API 層面，pytorch 爲我們提供了 torch.distributed.launch 啓動器，用於在命令行分佈式地執行 python 文件。在執行過程中，啓動器會將當前進程的（其實就是 GPU的）index 通過參數傳遞給 python，我們可以這樣獲得當前進程的 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 設置GPU 之間通信使用的後端和端口：

dist.init_process_group(backend='nccl')

之後，使用 DistributedSampler 對數據集進行劃分。如此前我們介紹的那樣，它能幫助我們將每個 batch 劃分成幾個 partition，在當前進程中只需要獲取和 rank 對應的那個 partition 進行訓練：

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

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

然後，使用 DistributedDataParallel 包裝模型，它能幫助我們爲不同 GPU 上求得的梯度進行 all reduce（即彙總不同 GPU 計算所得的梯度，並同步計算結果）。all reduce 後不同 GPU 中模型的梯度均爲 all reduce 之前各 GPU 梯度的均值：

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

最後，把數據和模型加載到當前進程使用的 GPU 中，正常進行正反向傳播：

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 並行訓練部分主要與如下代碼段有關：

# 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()

在使用時，調用 torch.distributed.launch 啓動器啓動：

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

在 ImageNet 上的完整訓練代碼，請點擊Github。

使用 torch.multiprocessing 取代啓動器

有的同學可能比較熟悉 torch.multiprocessing，也可以手動使用 torch.multiprocessing 進行多進程控制。繞開 torch.distributed.launch 自動控制開啓和退出進程的一些小毛病～

使用時，只需要調用 torch.multiprocessing.spawn，torch.multiprocessing 就會幫助我們自動創建進程。如下面的代碼所示，spawn 開啓了 nprocs=4 個線程，每個線程執行 main_worker 並向其中傳入 local_rank（當前進程 index）和 args（即 4 和 myargs）作爲參數：

import torch.multiprocessing as mp

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

這裏，我們直接將原本需要 torch.distributed.launch 管理的執行內容，封裝進 main_worker 函數中，其中 proc 對應 local_rank（當前進程 index），ngpus_per_node 對應 4， args 對應 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 讀取的默認環境變量作爲配置，我們需要手動爲 init_process_group 指定參數：

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

彙總一下，添加 multiprocessing 後並行訓練部分主要與如下代碼段有關：

# 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()

在使用時，直接使用 python 運行就可以了：

python main.py

在 ImageNet 上的完整訓練代碼，請點擊Github。

使用 Apex 再加速

Apex 是 NVIDIA 開源的用於混合精度訓練和分佈式訓練庫。Apex 對混合精度訓練的過程進行了封裝，改兩三行配置就可以進行混合精度的訓練，從而大幅度降低顯存佔用，節約運算時間。此外，Apex 也提供了對分佈式訓練的封裝，針對 NVIDIA 的 NCCL 通信庫進行了優化。

在混合精度訓練上，Apex 的封裝十分優雅。直接使用 amp.initialize 包裝模型和優化器，apex 就會自動幫助我們管理模型參數和優化器的精度了，根據精度需求不同可以傳入其他配置參數。

from apex import amp

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

在分佈式訓練的封裝上，Apex 在膠水層的改動並不大，主要是優化了 NCCL 的通信。因此，大部分代碼仍與 torch.distributed 保持一致。使用的時候只需要將 torch.nn.parallel.DistributedDataParallel 替換爲 apex.parallel.DistributedDataParallel 用於包裝模型。在 API 層面，相對於 torch.distributed ，它可以自動管理一些參數（可以少傳一點）：

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)

在正向傳播計算 loss 時，Apex 需要使用 amp.scale_loss 包裝，用於根據 loss 值自動對精度進行縮放：

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

彙總一下，Apex 的並行訓練部分主要與如下代碼段有關：

# 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()

在使用時，調用 torch.distributed.launch 啓動器啓動：

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

在 ImageNet 上的完整訓練代碼，請點擊Github。

Horovod 的優雅實現

Horovod 是 Uber 開源的深度學習工具，它的發展吸取了 Facebook "Training ImageNet In 1 Hour" 與百度 "Ring Allreduce" 的優點，可以無痛與 PyTorch/Tensorflow 等深度學習框架結合，實現並行訓練。

在 API 層面，Horovod 和 torch.distributed 十分相似。在 mpirun 的基礎上，Horovod 提供了自己封裝的 horovodrun 作爲啓動器。

與 torch.distributed.launch 相似，我們只需要編寫一份代碼，horovodrun 啓動器就會自動將其分配給個進程，分別在個 GPU 上運行。在執行過程中，啓動器會將當前進程的（其實就是 GPU的）index 注入 hvd，我們可以這樣獲得當前進程的 index：

import horovod.torch as hvd

hvd.local_rank()

與 init_process_group 相似，Horovod 使用 init 設置GPU 之間通信使用的後端和端口:

hvd.init()

接着，使用 DistributedSampler 對數據集進行劃分。如此前我們介紹的那樣，它能幫助我們將每個 batch 劃分成幾個 partition，在當前進程中只需要獲取和 rank 對應的那個 partition 進行訓練：

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

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

之後，使用 broadcast_parameters 包裝模型參數，將模型參數從編號爲 root_rank 的 GPU 複製到所有其他 GPU 中：

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

然後，使用 DistributedOptimizer 包裝優化器。它能幫助我們爲不同 GPU 上求得的梯度進行 all reduce（即彙總不同 GPU 計算所得的梯度，並同步計算結果）。all reduce 後不同 GPU 中模型的梯度均爲 all reduce 之前各 GPU 梯度的均值：

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

最後，把數據加載到當前 GPU 中。在編寫代碼時，我們只需要關注正常進行正向傳播和反向傳播：

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 的並行訓練部分主要與如下代碼段有關：

# 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()

在使用時，調用 horovodrun 啓動器啓動：

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

在 ImageNet 上的完整訓練代碼，請點擊Github。

尾註：

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

圖 2：配置詳情

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

圖 3：配置詳情

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

圖 4：配置詳情

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

https://zhuanlan.zhihu.com/p/98535650

並行訓練方法

Take-Away

簡單方便的 nn.DataParallel

使用 torch.distributed 加速並行訓練

使用 torch.multiprocessing 取代啓動器

使用 Apex 再加速

Horovod 的優雅實現

尾註：

redis的key亂碼問題和值自增問題

CORS error 但是 status code 是200 OK

一個開源且全面的C#算法實戰教程

一款.NET開源、功能強大、跨平臺的繪圖庫 - OxyPlot

壓縮上傳的GPU數據的方案

使用skopeo同步鏡像

上採樣和PixelShuffle

GAE(Generalized Advantage Estimation) PPO

CMA-ES算法流程

TRPO、PPO

讀論文Trust Region Policy Optimization

Mac下配置sublime實現LaTeX

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結