Pytorch的模型加速方法:Dataparallel (DP) 和 DataparallelDistributedparallel (DDP)

ZhiboZhao發表於2021-07-16

Dataparallel 和 DataparallelDistributed 的區別

一、Dataparallel(DP)

1.1 Dartaparallel 的使用方式

Dataparallel 的使用方式比較簡單,只需要一句話即可: net = nn.Dataparallel(net, device_ids, output_device)

其中,net 就是自己定義的網路例項,device_ids就是需要使用的顯示卡列表,output_device 表示引數輸出結果的裝置,預設情況下 output_device = device_ids[0]。因此在使用時經常發現第一塊卡所佔用的視訊記憶體會多一些。

1.2 Dataparallel 的基本原理

Dataparallel是資料分離型,其具體做法是:在前向傳播過程中,輸入資料會被分成多個子部分送到不同的 device 中進行計算,而網路模型則是在每個 device 上都拷貝一份,即:輸入的 batch 是平均分配到每個 device 中去,而網路模型需要拷貝到每個 device 中。在反向傳播過程中,每個副本積累的梯度會被累加到原始模組中,未指明 output_device 的情況下會在 device_ids[0] 上進行運算,更新好以後把權重分發到其餘卡。

1.3 Dataparallel 的注意事項

執行DataParallel模組之前,並行化模組必須在device_ids [0]上具有其引數和緩衝區。在執行DataParallel之前,會首先把其模型的引數放在device_ids[0]上。舉個例子,伺服器是八卡的伺服器,剛好前面序號是0的卡被別人佔用著,於是你只能用其他的卡來,比如你用2和3號卡,如果你直接指定 device_ids=[2, 3] 的話會出現模型初始化錯誤,類似於module沒有複製到在 device_ids[0] 上去。那麼你需要在執行train之前需要新增如下兩句話指定程式可見的devices,如下:

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

當新增這兩行程式碼後,那麼 device_ids[0] 預設的就是第2號卡,你的模型也會初始化在第2號卡上了,而不會佔用第0號卡了。設定上面兩行程式碼後,那麼對這個程式而言可見的只有2和3號卡,和其他的卡沒有關係,這是物理上的號卡,邏輯上來說其實是對應0和1號卡,即 device_ids[0] 對應的就是第2號卡,device_ids[1] 對應的就是第3號卡。(當然你要保證上面這兩行程式碼需要定義在下面兩行程式碼之前:

device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

1.4 Dataparallel 的優缺點

Dataparallel 的優點就是使用起來非常簡單,能夠使用多卡的視訊記憶體來處理資料。然而其缺點是:會造成負載不均衡的情況,成為限制模型訓練速度的瓶頸。

二、DataparallelDistributed(DDP)

2.1 DDP 的基本原理

DataparallelDistributed 在每次迭代中,作業系統會為每個GPU建立一個程式,每個程式具有自己的 optimizer ,並獨立完成所有的優化步驟,程式內與一般的訓練無異。在各程式梯度計算完成之後,各程式需要將梯度進行彙總平均,然後再由 rank=0 的程式,將其 broadcast 到所有程式。各程式用該梯度來更新引數。由於各程式中的模型,初始引數一致 (初始時刻進行一次 broadcast),而每次用於更新引數的梯度也一致,因此,各程式的模型引數始終保持一致。而在 DataParallel 中,全程維護一個 optimizer,對各 GPU 上梯度進行求和,而在主 GPU 進行引數更新,之後再將模型引數 broadcast 到其他 GPU。相較於 DataParalleltorch.distributed 傳輸的資料量更少,因此速度更快,效率更高。

2.2 DDP的使用方式

DDP使用起來比DP要麻煩一些,具體想要了解其中原理的可以參考下面幾篇文章:

https://blog.csdn.net/laizi_laizi/article/details/115299263

DataParallel & DistributedDataParallel分散式訓練 - 知乎 (zhihu.com)

最後,參考上述文章,整理出來了下面一份可以直接跑的程式碼,由於個人環境不同,可能在個別環境出現不適配的情況,可以參考上述文章進行修改。

################
## main.py檔案
import argparse
from tqdm import tqdm
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
# 新增:
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

### 1. 基礎模組 ### 
# 假設我們的模型是這個,與DDP無關
class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
# 假設我們的資料是這個
def get_dataset():
    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True, 
        download=True, transform=transform)
    # DDP:使用DistributedSampler,DDP幫我們把細節都封裝起來了。
    #      用,就完事兒!sampler的原理,第二篇中有介紹。
    train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)
    # DDP:需要注意的是,這裡的batch_size指的是每個程式下的batch_size。
    #      也就是說,總batch_size是這裡的batch_size再乘以並行數(world_size)。
    trainloader = torch.utils.data.DataLoader(my_trainset, 
        batch_size=16, num_workers=2, sampler=train_sampler)
    return trainloader
    
### 2. 初始化我們的模型、資料、各種配置  ####
# DDP:從外部得到local_rank引數
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank

# DDP:DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')  # nccl是GPU裝置上最快、最推薦的後端

# 準備資料,要在DDP初始化之後進行
trainloader = get_dataset()

# 構造模型
model = ToyModel().to(local_rank)
# DDP: Load模型要在構造DDP模型之前,且只需要在master上載入就行了。
ckpt_path = None
if dist.get_rank() == 0 and ckpt_path is not None:
    model.load_state_dict(torch.load(ckpt_path))
# DDP: 構造DDP model
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

# DDP: 要在構造DDP model之後,才能用model初始化optimizer。
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# 假設我們的loss是這個
loss_func = nn.CrossEntropyLoss().to(local_rank)

### 3. 網路訓練  ###
model.train()
iterator = tqdm(range(100))
for epoch in iterator:
    # DDP:設定sampler的epoch,
    # DistributedSampler需要這個來指定shuffle方式,
    # 通過維持各個程式之間的相同隨機數種子使不同程式能獲得同樣的shuffle效果。
    trainloader.sampler.set_epoch(epoch)
    # 後面這部分,則與原來完全一致了。
    for data, label in trainloader:
        data, label = data.to(local_rank), label.to(local_rank)
        optimizer.zero_grad()
        prediction = model(data)
        loss = loss_func(prediction, label)
        loss.backward()
        iterator.desc = "loss = %0.3f" % loss
        optimizer.step()
    # DDP:
    # 1. save模型的時候,和DP模式一樣,有一個需要注意的點:儲存的是model.module而不是model。
    #    因為model其實是DDP model,引數是被`model=DDP(model)`包起來的。
    # 2. 只需要在程式0上儲存一次就行了,避免多次儲存重複的東西。
    if dist.get_rank() == 0:
        torch.save(model.module.state_dict(), "%d.ckpt" % epoch)


################
## Bash執行
# DDP: 使用torch.distributed.launch啟動DDP模式
# 使用CUDA_VISIBLE_DEVICES,來決定使用哪些GPU
# CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 main.py

三、總結

總之Dataparellel和Distribution都是模型訓練加速的一種方法。Dataparallel (支援單機多卡),但是速度慢(主要原因是它採用parameter server 模式,一張主卡作為reducer,負載不均衡,主卡成為訓練瓶頸),在主GPU上進行梯度計算和更新,再將引數給其他gpu。而DDP則使用多執行緒進行加速,訓練速度得到了明顯的提升,但是程式碼修改起來比較麻煩,需要不斷試錯積累經驗。

相關文章