1. Take-Away
筆者使用 PyTorch 編寫了不同加速庫在 ImageNet 上的使用示例(單機多卡)。需要的同學可以當作 quickstart 將所需要的部分 copy 到自己的項目中(Github 請點擊下面鏈接):
nn.DataParallel[1] 簡單方便的 nn.DataParallel
torch.distributed[2] 使用 torch.distributed 加速并行訓練
torch.multiprocessing[3] 使用 torch.multiprocessing 取代啟動器
apex[4] 使用 apex 再加速
horovod[5] horovod 的優雅實現
注:分布式 evaluation[6]
這里,筆者記錄了使用 4 塊 Tesla V100-PICE 在 ImageNet 進行了運行時間的測試。
測試結果發現 Apex 的加速效果最好,但與 Horovod/Distributed 差別不大,平時可以直接使用內置的 Distributed。Dataparallel 較慢,不推薦使用。
(后續會補上 V100/K80 上的測試結果,穿插了一些試驗所以中斷了。)
簡要記錄一下不同庫的分布式訓練方式,當作代碼的 README(我真是個小機靈鬼)~
2. 簡單方便的 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[7]。
3. 使用 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='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[8]。
4. 使用 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),進程數 nprocs 對應 4, args 對應 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()
在上面的代碼中值得注意的是,由于沒有 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, 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()
在使用時,直接使用 python 運行就可以了:
python main.py
在 ImageNet 上的完整訓練代碼,請點擊Github[9]。
5.使用 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[10]。
6.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:
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[11]。
7.分布式 evaluation
all_reduce, barrier 等 API 是 distributed 中更為基礎和底層的 API。這些 API 可以幫助我們控制進程之間的交互,控制 GPU 數據的傳輸。在自定義 GPU 協作邏輯,匯總 GPU 間少量的統計信息時,大有用處。熟練掌握這些 API 也可以幫助我們自己設計、優化分布式訓練、測試流程。
最近,不少同學私信了我這樣的問題,
訓練樣本被切分成了若干個部分,被若干個進程分別控制運行在若干個 GPU 上,如何在進程間進行通信匯總這些(GPU 上的)信息?
使用一張卡進行推理、測試太慢了,如何使用 Distributed 進行分布式地推理和測試,并將結果匯總在一起?
......
要解決這些問題,我們需要一個更為基礎的 API,匯總記錄不同 GPU 上生成的準確率、損失函數等指標信息。這個 API 就是 torch.distributed.all_reduce。
圖2:all_reduce 示意圖
如上圖所示,它的工作過程包含以下三步:
在調用 all_reduce(tensor, op=...)后,當前進程會向其他進程發送 tensor(例如 rank 0 會發送 rank 0 的 tensor 到 rank 1、2、3)
同時,當前進程接受其他進程發來的 tensor(例如 rank 0 會接收 rank 1 的 tensor、rank 2 的 tensor、rank 3 的 tensor)。
在全部接收完成后,當前進程(例如rank 0)會對當前進程的和接收到的 tensor (例如 rank 0 的 tensor、rank 1 的 tensor、rank 2 的 tensor、rank 3 的 tensor)進行 op (例如求和)操作。
使用 torch.distributed.all_reduce(loss, op=torch.distributed.reduce_op.SUM),我們就能夠對不同數據切片(不同 GPU 上的訓練數據)的損失函數進行求和了。接著,我們只要再將其除以進程(GPU)數量 world_size就可以得到損失函數的平均值。正確率也能夠通過同樣方法進行計算:
# 原始代碼 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 中數據切片的統計信息,用于分布式的 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))
值得注意的是,為了同步各進程的計算進度,我們在 reduce 之前插入了一個同步 API torch.distributed.barrier()。在所有進程運行到這一步之前,先完成此前代碼的進程會等待其他進程。這使得我們能夠得到準確、有序的輸出。在 Horovod 中,我們無法使用 torch.distributed.barrier(),取而代之的是,我們可以在 allreduce 過程中指明:
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 上的完整訓練代碼,請點擊Github[12]。
8.尾注:
本文中使用的 V100-PICE (前 4 個 GPU)的配置:
圖 3:配置詳情
本文中使用的 V100 (前 4 個 GPU)的配置:
圖 4:配置詳情
本文中使用的 K80 (前 4 個 GPU)的配置:
圖 5:配置詳情
筆者本身是 CV 研究生,今天摸魚的時候一時興起研究了一下,后面再慢慢完善~
-
數據
+關注
關注
8文章
7006瀏覽量
88947 -
gpu
+關注
關注
28文章
4729瀏覽量
128897 -
模型
+關注
關注
1文章
3229瀏覽量
48811
原文標題:8.尾注:
文章出處:【微信號:zenRRan,微信公眾號:深度學習自然語言處理】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論