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

ZhiboZhao發表於2021-07-16
PyTorch模型Parallel

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。相較於 DataParallel，torch.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則使用多執行緒進行加速，訓練速度得到了明顯的提升，但是程式碼修改起來比較麻煩，需要不斷試錯積累經驗。

相關文章