1. 首頁 > 生活百科 > >

使用Pytorch進行多卡訓練

生活百科

當一塊GPU不夠用時，我們就需要使用多卡進行并行訓練 。其中多卡并行可分為數據并行和模型并行 。具體區別如下圖所示：

使用Pytorch進行多卡訓練

文章插圖
由于模型并行比較少用，這里只對數據并行進行記錄 。對于pytorch，有兩種方式可以進行數據并行：數據并行（DataParallel, DP）和分布式數據并行（DistributedDataParallel, DDP） 。
在多卡訓練的實現上，DP與DDP的思路是相似的：
1、每張卡都復制一個有相同參數的模型副本 。
2、每次迭代，每張卡分別輸入不同批次數據，分別計算梯度 。
3、DP與DDP的主要不同在于接下來的多卡通信：
DP的多卡交互實現在一個進程之中，它將一張卡視為主卡，維護單獨模型優化器 。所有卡計算完梯度后，主卡匯聚其它卡的梯度進行平均并用優化器更新模型參數，再將模型參數更新至其它卡上 。
DDP則分別為每張卡創建一個進程，每個進程相應的卡上都獨立維護模型和優化器 。在每次每張卡計算完梯度之后，進程之間以NCLL（NVIDIA GPU通信）為通信后端，使各卡獲取其它卡的梯度 。各卡對獲取的梯度進行平均 ， 然后執行后續的參數更新 。由于每張卡上的模型與優化器參數在初始化時就保持一致 ， 而每次迭代的平均梯度也保持一致，那么即使沒有進行參數復制，所有卡的模型參數也是保持一致的 。
Pytorch官方推薦我們使用DDP 。DP經過我的實驗，兩塊GPU甚至比一塊還慢 。當然不同模型可能有不同的結果 。下面分別對DP和DDP進行記錄 。
DPPytorch的DP實現多GPU訓練十分簡單 ， 只需在單GPU的基礎上加一行代碼即可 。以下是一個DEMO的代碼 。
import torchfrom torch import nnfrom torch.optim import Adamfrom torch.nn.parallel import DataParallelclass DEMO_model(nn.Module):def __init__(self, in_size, out_size):super().__init__()self.fc = nn.Linear(in_size, out_size)def forward(self, inp):outp = self.fc(inp)print(inp.shape, outp.device)return outpmodel = DEMO_model(10, 5).to('cuda')model = DataParallel(model, device_ids=[0, 1]) # 額外加這一行adam = Adam(model.parameters())# 進行訓練for i in range(1):x = torch.rand([128, 10]) # 獲取訓練數據，無需指定設備y = model(x) # 自動均勻劃分數據批量并分配至各GPU ， 輸出結果y會聚集到GPU0中loss = torch.norm(y)loss.backward()adam.step()其中model = DataParallel(model, device_ids=[0, 1])這行將模型復制到0,1號GPU上 。輸入數據x無需指定設備 ， 它將會被均勻分配至各塊GPU模型 ， 進行前向傳播 。之后各塊GPU的輸出再合并到GPU0中，得到輸出y 。輸出y在GPU0中計算損失，并進行反向傳播計算梯度、優化器更新參數 。
DDP為了對分布式編程有基本概念，首先使用pytorch內部的方法實現一個多進程程序，再使用DDP模塊實現模型的分布式訓練 。
Pytorch分布式基礎首先使用pytorch內部的方法編寫一個多進程程序作為編寫分布式訓練的基礎 。
import os, torchimport torch.multiprocessing as mpimport torch.distributed as distdef run(rank, size):tensor = torch.tensor([1,2,3,4], device='cuda:'+str(rank)) # ——1——group = dist.new_group(range(size)) # ——2——dist.all_reduce(tensor=tensor, group=group, op=dist.ReduceOp.SUM) # ——3——print(str(rank)+ ': ' + str(tensor) + '\n')def ini_process(rank, size, fn, backend = 'nccl'):os.environ['MASTER_ADDR'] = '127.0.0.1' # ——4——os.environ['MASTER_PORT'] = '1234'dist.init_process_group(backend, rank=rank, world_size=size) # ——5——fn(rank, size) # ——6——if __name__ == '__main__': # ——7——mp.set_start_method('spawn') # ——8——size = 2 # ——9——ps = []for rank in range(size):p = mp.Process(target=ini_process, args=(rank, size, run)) # ——10——p.start()ps.append(p)for p in ps: # ——11——p.join()以上代碼主進程創建了兩個子進程，子進程之間使用NCCL后端進行通信 。每個子進程各占用一個GPU資源，實現了所有GPU張量求和的功能 。細節注釋如下：
1、為每個子進程定義相同名稱的張量，并分別分配至不同的GPU ， 從而能進行后續的GPU間通信 。
2、定義一個通信組，用于后面的all_reduce通信操作 。
3、all_reduce操作以及其它通信方式請看下圖：
使用Pytorch進行多卡訓練

文章插圖
4、定義編號(rank)為0的ip和端口地址，讓每個子進程都知道 。ip和端口地址可以隨意定義 ， 不沖突即可 。如果不設置，子進程在涉及進程通信時會出錯 。

推薦閱讀