pytorch的多GPU訓(xùn)練的方式有哪些

這篇文章將為大家詳細(xì)講解有關(guān)pytorch的多GPU訓(xùn)練的方式有哪些，小編覺得挺實(shí)用的，因此分享給大家做個(gè)參考，希望大家閱讀完這篇文章后可以有所收獲。

方法一：torch.nn.DataParallel

1. 原理

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

這里的作業(yè)就是pytorch中要處理的data。

與此同時(shí)，他也可以先花3min把作業(yè)分配給3個(gè)同伙，大家一起60min做完。最后他再花3min把作業(yè)收起來，一共需要66min。

這個(gè)小朋友就是主GPU。他的過程是：分發(fā) ->并行運(yùn)算->結(jié)果回收。 

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

這種方法也稱為單進(jìn)程多GPU訓(xùn)練模式:DP模式，這種并行模式下并行的多卡都是由一個(gè)進(jìn)程進(jìn)行控制。換句話說，在進(jìn)行梯度的傳播時(shí)，是在主GPU上進(jìn)行的。

采用torch.nn.DataParallel進(jìn)行多GPU并行訓(xùn)練時(shí)，與其搭配的數(shù)據(jù)讀取代碼是：torch.utils.data.DataLoader

2. 常用的配套代碼如下

train_datasets = customData(train_txt)  #創(chuàng)建datasettrain_dataloaders = torch.utils.data.DataLoader(train_datasets,opt.batch_size,num_workers=train_num_workers,shuffle=True)  #創(chuàng)建dataloadermodel = efficientnet_b0(num_classes = opt.num_class)  #創(chuàng)建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，也就是分發(fā)任務(wù)和結(jié)果回收的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))  #將數(shù)據(jù)放到主要GPU   labels = Variable(labels.to(device))

3. 優(yōu)缺點(diǎn)

• 優(yōu)點(diǎn)：配置起來非常方便

• 缺點(diǎn)：GPU負(fù)載不均衡，主GPU的負(fù)載很大，而其他GPU的負(fù)載很少

方法二：torch.distributed

1. 代碼說明

這個(gè)方法本來是用于多機(jī)器多卡（多節(jié)點(diǎn)多卡）訓(xùn)練的，但是也可以用于單機(jī)多卡（即將節(jié)點(diǎn)數(shù)設(shè)置為1）訓(xùn)練。

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

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

這里給出一個(gè)簡(jiǎn)單的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） 配置每個(gè)進(jìn)程的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) 封裝之前要把模型移到對(duì)應(yīng)的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）啟動(dòng)方式：在torch.distributed當(dāng)中提供了一個(gè)用于啟動(dòng)的程序torch.distributed.launch，此幫助程序可用于為每個(gè)節(jié)點(diǎn)啟動(dòng)多個(gè)進(jìn)程以進(jìn)行分布式訓(xùn)練，它在每個(gè)訓(xùn)練節(jié)點(diǎn)上產(chǎn)生多個(gè)分布式訓(xùn)練進(jìn)程。

（2）啟動(dòng)命令：

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

這里需要說明一下參數(shù)：

• CUDA_VISIBLE_DEVICES：設(shè)置我們可用的GPU的id

• torch.distributed.launch：用于啟動(dòng)多節(jié)點(diǎn)多GPU的訓(xùn)練

• nproc_per_node：表示設(shè)置的進(jìn)程數(shù)量，一般情況設(shè)置為可用的GPU數(shù)量，即有多少個(gè)可用的GPU就設(shè)置多少個(gè)進(jìn)程。

• local rank：關(guān)于這個(gè)參數(shù)的意義，我們將在后面的情形中進(jìn)行說明。

（3）一些情形的說明：

情形1：直接運(yùn)行上述的命令

運(yùn)行的結(jié)果如下：

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，其數(shù)量與我們?cè)O(shè)置的nproc_per_node是一樣的，與我們?cè)O(shè)置的可用GPU的數(shù)量是無關(guān)的。這里就要說明一下local rank的意義。

local rank：表示的是當(dāng)前的進(jìn)程在當(dāng)前節(jié)點(diǎn)的編號(hào)，因?yàn)槲覀冊(cè)O(shè)置了2個(gè)進(jìn)程，因此進(jìn)程的編號(hào)就是0和1

在很多博客中都直接說明local_rank等于進(jìn)程內(nèi)的GPU編號(hào)，這種說法實(shí)際上是不準(zhǔn)確的。這個(gè)編號(hào)并不是GPU的編號(hào)??！

在使用啟動(dòng)命令時(shí)，torch.distributed.launch工具會(huì)默認(rèn)地根據(jù)nproc_per_node傳入local_rank參數(shù)，之后再通過下面的代碼可以得到local_rank.

local_rank = torch.distributed.get_rank()

因?yàn)槭悄J(rèn)傳入?yún)?shù)local_rank，所以還可以這么寫，其輸出與torch.distributed.get_rank()相同

import argparse
parser = argparse.ArgumentParser()
# 注意這個(gè)參數(shù)，必須要以這種形式指定，即使代碼中不使用。因?yàn)?nbsp;launch 工具默認(rèn)傳遞該參數(shù)
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設(shè)置為4，即將進(jìn)程數(shù)設(shè)置為可用的GPU數(shù)

運(yùn)行結(jié)果如下：

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])

可以看到，此時(shí)的local_rank共有4個(gè)，與進(jìn)程數(shù)相同。并且我們?cè)O(shè)置的可用GPU的id是1，2，3，4，而local_rank的輸出為0，1，2，3，可見local_rank并不是GPU的編號(hào)。

雖然在代碼中模型并行的device_ids設(shè)置為local_rank，而local_rank為0，1，2，3，但是實(shí)際上還是采用可用的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設(shè)置為4，但是不設(shè)置可用的GPU ID

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

此時(shí)我們?cè)偈褂胣vidia-smi來查看GPU的使用情況，如下?？梢钥吹酱藭r(shí)使用的GPU就是local rank的id。相比于情形2，我們可以總結(jié)：

當(dāng)沒有設(shè)置可用的GPU ID時(shí)，所采用的GPU id就等于local rank的id。本質(zhì)上是將進(jìn)程的編號(hào)作為GPU編號(hào)使用，因此local_rank等于進(jìn)程的編號(hào)這個(gè)定義是不變的。

當(dāng)設(shè)置可用的GPU ID，所采用的GPU id就等于GPU id。

情形4：將nproc_per_node設(shè)置為5，即超出了可以用的GPU數(shù)

輸出結(jié)果如下，可以看到是報(bào)錯(cuò)的，因?yàn)檫M(jìn)程數(shù)超出了可以用的GPU數(shù)量

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

關(guān)于“pytorch的多GPU訓(xùn)練的方式有哪些”這篇文章就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，使各位可以學(xué)到更多知識(shí)，如果覺得文章不錯(cuò)，請(qǐng)把它分享出去讓更多的人看到。