PyTorch?Distributed?Data?Parallel如何使用

這篇文章主要介紹“PyTorch Distributed Data Parallel如何使用”，在日常操作中，相信很多人在PyTorch Distributed Data Parallel如何使用問題上存在疑惑，小編查閱了各式資料，整理出簡單好用的操作方法，希望對大家解答”PyTorch Distributed Data Parallel如何使用”的疑惑有所幫助！接下來，請跟著小編一起來學習吧！

DDP

Distributed Data Parallel 簡稱 DDP，是 PyTorch 框架下一種適用于單機多卡、多機多卡任務的數據并行方式。由于其良好的執行效率及廣泛的顯卡支持，熟練掌握 DDP 已經成為深度學習從業者所必備的技能之一。

概念辨析

具體講解 DDP 之前，我們先了解了解它和 Data Parallel (DP) 之間的區別。DP 同樣是 PyTorch 常見的多 GPU 并行方式之一，且它的實現非常簡潔：

# 函數定義
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
'''
module : 模型
device_ids : 參與訓練的 GPU 列表
output_device : 指定輸出的 GPU, 通常省略, 即默認使用索引為 0 的顯卡
'''
# 程序模板
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

基本原理及固有缺陷：在 Data Parallel 模式下，數據會被自動切分，加載到 GPU。同時，模型也將拷貝至各個 GPU 進行正向傳播。在多個進程之間，會有一個進程充當 master 節點，負責收集各張顯卡積累的梯度，并據此更新參數，再統一發送至其他顯卡。因此整體而言，master 節點承擔了更多的計算與通信任務，容易造成網絡堵塞，影響訓練速度。

常見問題及解決方案：Data Parallel 要求模型必須在 device_ids[0] 擁有參數及緩沖區，因此當卡 0 被占用時，可以在 nn.DataParallel 之前添加如下代碼：

# 按照 PIC_BUS_ID 順序自 0 開始排列 GPU 設備
os.environ['CUDA_DEVICE_ORDER'] = 'PIC_BUS_ID'
# 設置當前使用的 GPU 為 2、3 號設備
os.environ['CUDA_VISIBLE_DEVICES'] = '2, 3'

如此，device_ids[0] 將被默認為 2 號卡，device_ids[1] 則對應 3 號卡

相較于 DP, Distributed Data Parallel 的實現要復雜得多，但是它的優勢也非常明顯：

• DDP 速度更快，可以達到略低于顯卡數量的加速比；

• DDP 可以實現負載的均勻分配，克服了 DP 需要一個進程充當 master 節點的固有缺陷；

• 采用 DDP 通常可以支持更大的 batch size，不會像 DP 那樣出現其他顯卡尚有余力，而卡 0 直接 out of memory 的情況；

• 另外，在 DDP 模式下，輸入到 data loader 的 bacth size 不再代表總數，而是每塊 GPU 各自負責的 sample 數量。比方說，batch_size = 30，有兩塊 GPU。在 DP 模式下，每塊 GPU 會負責 15 個樣本。而在 DDP 模式下，每塊 GPU 會各自負責 30 個樣本；

• DDP 基本原理：倘若我們擁有 N 張顯卡，則在 Distributed Data Parallel 模式下，就會啟動 N 個進程。每個進程在各自的卡上加載模型，且模型的參數完全相同。訓練過程中，各個進程通過一種名為 Ring-Reduce 的方式與其他進程通信，交換彼此的梯度，從而獲得所有的梯度信息。隨后，各個進程利用梯度的平均值更新參數。由于初始值和更新量完全相同，所以各個進程更新后的參數仍保持一致。

常用術語

• rank

• 進程號

• 多進程上下文中，通常假定 rank = 0 為主進程或第一個進程

• node

• 物理節點，表示一個容器或一臺機器

• 節點內部可以包含多個 GPU

• local_rank

• 一個 node 中，進程的相對序號

• local_rank 在 node 之間獨立

• world_size

• 全局進程數

• 一個分布式任務中 rank 的數量

• group

• 進程組

• 一個分布式任務就對應一個進程組

• 只有當用戶創立多個進程組時，才會用到

代碼實現

Distributed Data Parallel 可以通過 Python 的 torch.distributed.launch 啟動器，在命令行分布式地執行 Python 文件。執行過程中，啟動器會將當前進程（其實就是 GPU）的 index 通過參數傳遞給 Python，而我們可以利用如下方式獲取這個 index：

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    metavar='N', help='Local process rank.')
args = parser.parse_args()
# print(args.local_rank)
# local_rank 表示本地進程序號

隨后，初始化進程組。對于在 GPU 執行的任務，建議選擇 nccl (由 NVIDIA 推出) 作為通信后端。對于在 CPU 執行的任務，建議選擇 gloo (由 Facebook 推出) 作為通信后端。倘若不傳入 init_method，則默認為 env://，表示自環境變量讀取分布式信息

dist.init_process_group(backend='nccl', init_method='env://')
# 初始化進程組之后, 通常會執行這兩行代碼
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)
# 后續的 model = model.to(device), tensor.cuda(device)
# 對應的都是這里由 args.local_rank 初始化得到的 device

數據部分，使用 Distributed Sampler 劃分數據集，并將 sampler 傳入 data loader。需要注意的是，此時在 data loader 中不能指定 shuffle 為 True，否則會報錯 (sampler 已具備隨機打亂功能)

dev_sampler = data.DistributedSampler(dev_data_set)
train_sampler = data.DistributedSampler(train_data_set)
dev_loader = data.DataLoader(dev_data_set, batch_size=dev_batch_size, 
                             shuffle=False, sampler=dev_sampler)
train_loader = data.DataLoader(train_data_set, batch_size=train_batch_size, 
                               shuffle=False, sampler=train_sampler)

模型部分，首先將將模型送至 device，即對應的 GPU 上，再使用 Distributed Data Parallel 包裝模型（順序顛倒會報錯）

model = model.to(device)
model = nn.parallel.DistributedDataParallel(
    model, device_ids=[args.local_rank], output_device=args.local_rank
)

Distributed Data Parallel 模式下，保存模型應使用 net.module.state_dict()，而非 net.state_dict()。且無論是保存模型，還是 LOGGER 打印，只對 local_rank 為 0 的進程操作即可，因此代碼中會有很多 args.local_rank == 0 的判斷

if args.local_rank == 0:
    LOGGER.info(f'saving latest model: {output_path}')
    torch.save({'model': model.module.state_dict(), 
                'optimizer': None, 'epoch': epoch, 'best-f1': best_f1}, 
               open(os.path.join(output_path, 'latest_model_{}.pth'.format(fold)), 'wb'))

利用 torch.load 加載模型時，設置 map_location=device，否則卡 0 會承擔更多的開銷

load_model = torch.load(best_path, map_location=device)
model.load_state_dict(load_model['model'])

• dist.barrier() 可用于同步多個進程，建議只在必要的位置使用，如初始化 DDP 模型之前、權重更新之后、開啟新一輪 epoch 之前

• 計算 accuracy 時，可以使用 dist.all_reduce(score, op=dist.ReduceOp.SUM)，將各個進程計算的準確率求平均

• 計算 f1-score 時，可以使用 dist.all_gather(all_prediction_list, prediction_list)，將各個進程獲得的預測值和真實值匯總到 all_list，再統一代入公式

啟動方式

torch.distributed.launch

# 此處 --nproc_per_node 4 的含義是 server 有 4 張顯卡
python torch.distributed.launch --nproc_per_node 4 train.py
# 倘若使用 nohup, 則注意輸入命令后 exit 當前終端
python torch.distributed.launch --nproc_per_node 4 train.py

• torchrun，推薦使用這種方式，因為 torch.distributed.launch 即將棄用

代碼中，只需將 Argument Parser 相關的部分替換為

local_rank = int(os.environ['LOCAL_RANK'])

然后將 args.local_rank 全部改為 local_rank 即可

啟動命令

# 單機多卡訓練時, 可以不指定 nnodes
torchrun --nnodes=1 --nproc_per_node=4 train.py
# 倘若使用 nohup, 則注意輸入命令后 exit 當前終端
nohup torchrun --nnodes=1 --nproc_per_node=4 train.py > nohup.out &

到此，關于“PyTorch Distributed Data Parallel如何使用”的學習就結束了，希望能夠解決大家的疑惑。理論與實踐的搭配能更好的幫助大家學習，快去試試吧！若想繼續學習更多相關知識，請繼續關注億速云網站，小編會繼續努力為大家帶來更多實用的文章！