詳解pytorch的多GPU訓練的兩種方式

Posted on by

方法一：torch.nn.DataParallel

1. 原理

如下圖所示：小朋友一個人做4份作業，假設1份需要60min，共需要240min。

這裡的作業就是pytorch中要處理的data。

與此同時，他也可以先花3min把作業分配給3個同夥，大傢一起60min做完。最後他再花3min把作業收起來，一共需要66min。

這個小朋友就是主GPU。他的過程是：分發 ->並行運算->結果回收。 

這就是pytorch要使用的第一種並行方法：torch.nn.DataParallel

這種方法也稱為單進程多GPU訓練模式:DP模式，這種並行模式下並行的多卡都是由一個進程進行控制。換句話說，在進行梯度的傳播時，是在主GPU上進行的。

采用torch.nn.DataParallel進行多GPU並行訓練時，與其搭配的數據讀取代碼是：torch.utils.data.DataLoader

2. 常用的配套代碼如下

train_datasets = customData(train_txt)  #創建datasettrain_dataloaders = torch.utils.data.DataLoader(train_datasets,opt.batch_size,num_workers=train_num_workers,shuffle=True)  #創建dataloadermodel = efficientnet_b0(num_classes = opt.num_class)  #創建modeldevice_list = list(map(int,list(opt.device_id)))print("Using gpu"," ".join([str(v) for v in device_list]))device = device_list[0]  #主GPU，也就是分發任務和結果回收的GPU，也是梯度傳播更新的GPUmodel = torch.nn.DataParallel(model,device_ids=device_list)model.to(device)for data in train_dataloaders:    model.train(True)   inputs, labels = data   inputs = Variable(inputs.to(device))  #將數據放到主要GPU   labels = Variable(labels.to(device))

3. 優缺點

  • 優點：配置起來非常方便
  • 缺點：GPU負載不均衡，主GPU的負載很大，而其他GPU的負載很少

方法二：torch.distributed

1. 代碼說明

這個方法本來是用於多機器多卡（多節點多卡）訓練的，但是也可以用於單機多卡（即將節點數設置為1）訓練。

初始化的代碼如下，這個一定要寫在最前面。

from torch.utils.data.distributed import DistributedSampler
torch.distributed.init_process_group(backend="nccl")

這裡給出一個簡單的demo.py作為說明：

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
import os
from torch.utils.data.distributed import DistributedSampler
# 1) 初始化
torch.distributed.init_process_group(backend="nccl")
 
input_size = 5
output_size = 2
batch_size = 30
data_size = 90
 
# 2） 配置每個進程的gpu
local_rank = torch.distributed.get_rank()
print('local_rank',local_rank)
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
 
class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size).to('cuda')
 
    def __getitem__(self, index):
        return self.data[index]
 
    def __len__(self):
        return self.len
 
dataset = RandomDataset(input_size, data_size)
# 3）使用DistributedSampler
rand_loader = DataLoader(dataset=dataset,
                         batch_size=batch_size,
                         sampler=DistributedSampler(dataset))
 
class Model(nn.Module):
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)
 
    def forward(self, input):
        output = self.fc(input)
        print("  In Model: input size", input.size(),
              "output size", output.size())
        return output
 
model = Model(input_size, output_size)
 
# 4) 封裝之前要把模型移到對應的gpu
model.to(device)
 
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # 5) 封裝
    model = torch.nn.parallel.DistributedDataParallel(model,
                                                      device_ids=[local_rank],
                                                      output_device=local_rank)
 
for data in rand_loader:
    if torch.cuda.is_available():
        input_var = data
    else:
        input_var = data
 
    output = model(input_var)
    print("Outside: input size", input_var.size(), "output_size", output.size())

（1）啟動方式：在torch.distributed當中提供瞭一個用於啟動的程序torch.distributed.launch，此幫助程序可用於為每個節點啟動多個進程以進行分佈式訓練，它在每個訓練節點上產生多個分佈式訓練進程。

（2）啟動命令：

CUDA_VISIBLE_DEVICES=1,2,3,4 python -m torch.distributed.launch --nproc_per_node=2 torch_ddp.py

這裡需要說明一下參數：

  • CUDA_VISIBLE_DEVICES：設置我們可用的GPU的id
  • torch.distributed.launch：用於啟動多節點多GPU的訓練
  • nproc_per_node：表示設置的進程數量，一般情況設置為可用的GPU數量，即有多少個可用的GPU就設置多少個進程。
  • local rank：關於這個參數的意義，我們將在後面的情形中進行說明。

（3）一些情形的說明：

情形1：直接運行上述的命令

運行的結果如下：

local_rank 1
local_rank 0
Let's use 4 GPUs!
Let's use 4 GPUs!
  In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside: input size torch.Size([15, 5]) output_size torch.Size([15, 2])
  In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2])
Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside: input size torch.Size([15, 5]) output_size torch.Size([15, 2])

可以看到local rank的輸出為0和1，其數量與我們設置的nproc_per_node是一樣的，與我們設置的可用GPU的數量是無關的。這裡就要說明一下local rank的意義。

local rank：表示的是當前的進程在當前節點的編號，因為我們設置瞭2個進程，因此進程的編號就是0和1

在很多博客中都直接說明local_rank等於進程內的GPU編號，這種說法實際上是不準確的。這個編號並不是GPU的編號！！

在使用啟動命令時，torch.distributed.launch工具會默認地根據nproc_per_node傳入local_rank參數，之後再通過下面的代碼可以得到local_rank.

local_rank = torch.distributed.get_rank()

因為是默認傳入參數local_rank，所以還可以這麼寫，其輸出與torch.distributed.get_rank()相同

import argparse
parser = argparse.ArgumentParser()
# 註意這個參數，必須要以這種形式指定，即使代碼中不使用。因為 launch 工具默認傳遞該參數
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
 
local_rank = args.local_rank
print('local_rank',args.local_rank)

 情形2：將nproc_per_node設置為4，即將進程數設置為可用的GPU數

運行結果如下：

local_rank 2
local_rank 3
local_rank 1
local_rank 0
Let's use 4 GPUs!
Let's use 4 GPUs!
Let's use 4 GPUs!
Let's use 4 GPUs!
  In Model: input size torch.Size([23, 5]) output size torch.Size([23, 2])
Outside: input size torch.Size([23, 5]) output_size torch.Size([23, 2])
  In Model: input size torch.Size([23, 5]) output size torch.Size([23, 2])
Outside: input size torch.Size([23, 5]) output_size torch.Size([23, 2])
  In Model: input size torch.Size([23, 5]) output size torch.Size([23, 2])
Outside: input size torch.Size([23, 5]) output_size torch.Size([23, 2])
  In Model: input size torch.Size([23, 5]) output size torch.Size([23, 2])
Outside: input size torch.Size([23, 5]) output_size torch.Size([23, 2])

可以看到，此時的local_rank共有4個，與進程數相同。並且我們設置的可用GPU的id是1，2，3，4，而local_rank的輸出為0，1，2，3，可見local_rank並不是GPU的編號。

雖然在代碼中模型並行的device_ids設置為local_rank，而local_rank為0，1，2，3，但是實際上還是采用可用的GPU：1，2，3，4。可以通過nvidia-smi來查看，PID為86478，86479，86480，864782。

model = torch.nn.parallel.DistributedDataParallel(model,
                                             device_ids=[local_rank],
                                             output_device=local_rank)

情形3：將nproc_per_node設置為4，但是不設置可用的GPU ID

python -m torch.distributed.launch --nproc_per_node=4 ddp.py

此時我們再使用nvidia-smi來查看GPU的使用情況，如下。可以看到此時使用的GPU就是local rank的id。相比於情形2，我們可以總結：

當沒有設置可用的GPU ID時，所采用的GPU id就等於local rank的id。本質上是將進程的編號作為GPU編號使用，因此local_rank等於進程的編號這個定義是不變的。

當設置可用的GPU ID，所采用的GPU id就等於GPU id。

情形4：將nproc_per_node設置為5，即超出瞭可以用的GPU數

輸出結果如下，可以看到是報錯的，因為進程數超出瞭可以用的GPU數量

local_rank 3
local_rank 2
local_rank 4
local_rank 1
local_rank 0
THCudaCheck FAIL file=/pytorch/torch/csrc/cuda/Module.cpp line=59 error=101 : invalid device ordinal
Traceback (most recent call last):
  File "ddp.py", line 18, in <module>
    torch.cuda.set_device(local_rank)
  File "/home/yckj3822/anaconda3/lib/python3.6/site-packages/torch/cuda/__init__.py", line 281, in set_device
    torch._C._cuda_setDevice(device)
RuntimeError: cuda runtime error (101) : invalid device ordinal at /pytorch/torch/csrc/cuda/Module.cpp:59

到此這篇關於詳解pytorch的多GPU訓練的兩種方式的文章就介紹到這瞭,更多相關pytorch的多GPU訓練內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet！

推薦閱讀：