書圈 - 必備的5種Pytorch並行訓練方法－鑽石舞台

點擊上圖，查看教學大綱

來源 | 極市平台

又到適宜划水的周五啦，機器在學習，人很無聊。在打開 b 站「學習」之前看着那空着一半的顯卡決定寫點什麼餵飽它們～因此，從 V100-PICE/V100/K80 中各拿出 4 張卡，試驗一下哪種分布式學習庫速度最快！這下終於能把剩下的顯存吃完啦，又是老師的勤奮好學生啦（我真是個小機靈鬼）!

Take-Away

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

1、簡單方便的 nn.DataParallel

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

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

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

3、使用 torch.multiprocessing 取代啟動器

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 的優雅實現

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

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

簡要記錄一下不同庫的分布式訓練方式，當作代碼的 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])forepochinrange(100):forbatch_idx,(data,target)inenumerate(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.pyimporttorchimporttorch.distributedasdistgpus=[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())forepochinrange(100):forbatch_idx,(data,target)inenumerate(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 執行即可：

pythonmain.py

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

使用torch.distributed加速並行訓練

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

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

在 API 層面，pytorch 為我們提供了 torch.distributed.launch 啟動器，用於在命令行分布式地執行 python 文件。在執行過程中，啟動器會將當前進程的（其實就是 GPU的）index 通過參數傳遞給 python，我們可以這樣獲得當前進程的 index：

parser=argparse.ArgumentParser()parser.add_argument('--local_rank',default=-1,type=int,help='noderankfordistributedtraining')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()forepochinrange(100):forbatch_idx,(data,target)inenumerate(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.pyimporttorchimportargparseimporttorch.distributedasdistparser=argparse.ArgumentParser()parser.add_argument('--local_rank',default=-1,type=int,help='noderankfordistributedtraining')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())forepochinrange(100):forbatch_idx,(data,target)inenumerate(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,3python-mtorch.distributed.launch--nproc_per_node=4main.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）作為參數：

importtorch.multiprocessingasmpmp.spawn(main_worker,nprocs=4,args=(4,myargs))

這裡，我們直接將原本需要 torch.distributed.launch 管理的執行內容，封裝進 main_worker 函數中，其中 proc 對應 local_rank（當前進程 index），ngpus_per_node 對應 4， args 對應 myargs：

defmain_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())forepochinrange(100):forbatch_idx,(data,target)inenumerate(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.pyimporttorchimporttorch.distributedasdistimporttorch.multiprocessingasmpmp.spawn(main_worker,nprocs=4,args=(4,myargs))defmain_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())forepochinrange(100):forbatch_idx,(data,target)inenumerate(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 運行就可以了：

pythonmain.py

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

使用Apex再加速

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

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

fromapeximportampmodel,optimizer=amp.initialize(model,optimizer

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

fromapex.parallelimportDistributedDataParallelmodel=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 值自動對精度進行縮放：

withamp.scale_loss(loss,optimizer)asscaled_loss:scaled_loss.backward()

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

#main.pyimporttorchimportargparseimporttorch.distributedasdistfromapex.parallelimportDistributedDataParallelparser=argparse.ArgumentParser()parser.add_argument('--local_rank',default=-1,type=int,help='noderankfordistributedtraining')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())forepochinrange(100):forbatch_idx,(data,target)inenumerate(train_loader):images=images.cuda(non_blocking=True)target=target.cuda(non_blocking=True)...output=model(images)loss=criterion(output,target)optimizer.zero_grad()withamp.scale_loss(loss,optimizer)asscaled_loss:scaled_loss.backward()optimizer.step()

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

UDA_VISIBLE_DEVICES=0,1,2,3python-mtorch.distributed.launch--nproc_per_node=4main.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 啟動器就會自動將其分配給n個進程，分別在n個 GPU 上運行。在執行過程中，啟動器會將當前進程的（其實就是 GPU的）index 注入 hvd，我們可以這樣獲得當前進程的 index：

importhorovod.torchashvdhvd.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)forepochinrange(100):forbatch_idx,(data,target)inenumerate(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.pyimporttorchimporthorovod.torchashvdhvd.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)forepochinrange(100):forbatch_idx,(data,target)inenumerate(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,3horovodrun-np4-Hlocalhost:4--verbosepythonmain.py

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

尾注

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