詳解pytorch的多GPU訓(xùn)練的兩種方式
方法一:torch.nn.DataParallel
1. 原理
如下圖所示:小朋友一個人做4份作業(yè),假設(shè)1份需要60min,共需要240min。

這里的作業(yè)就是pytorch中要處理的data。
與此同時,他也可以先花3min把作業(yè)分配給3個同伙,大家一起60min做完。最后他再花3min把作業(yè)收起來,一共需要66min。

這個小朋友就是主GPU。他的過程是:分發(fā) ->并行運(yùn)算->結(jié)果回收。
這就是pytorch要使用的第一種并行方法:torch.nn.DataParallel
這種方法也稱為單進(jìn)程多GPU訓(xùn)練模式:DP模式,這種并行模式下并行的多卡都是由一個進(jìn)程進(jìn)行控制。換句話說,在進(jìn)行梯度的傳播時,是在主GPU上進(jìn)行的。

采用torch.nn.DataParallel進(jìn)行多GPU并行訓(xùn)練時,與其搭配的數(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. 代碼說明
這個方法本來是用于多機(jī)器多卡(多節(jié)點(diǎn)多卡)訓(xùn)練的,但是也可以用于單機(jī)多卡(即將節(jié)點(diǎn)數(shù)設(shè)置為1)訓(xùn)練。
初始化的代碼如下,這個一定要寫在最前面。
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) 配置每個進(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) 封裝之前要把模型移到對應(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)啟動方式:在torch.distributed當(dāng)中提供了一個用于啟動的程序torch.distributed.launch,此幫助程序可用于為每個節(jié)點(diǎn)啟動多個進(jìn)程以進(jìn)行分布式訓(xùn)練,它在每個訓(xùn)練節(jié)點(diǎn)上產(chǎn)生多個分布式訓(xùn)練進(jìn)程。
(2)啟動命令:
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:用于啟動多節(jié)點(diǎn)多GPU的訓(xùn)練
- nproc_per_node:表示設(shè)置的進(jìn)程數(shù)量,一般情況設(shè)置為可用的GPU數(shù)量,即有多少個可用的GPU就設(shè)置多少個進(jìn)程。
- local rank:關(guān)于這個參數(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ù)量與我們設(shè)置的nproc_per_node是一樣的,與我們設(shè)置的可用GPU的數(shù)量是無關(guān)的。這里就要說明一下local rank的意義。
local rank:表示的是當(dāng)前的進(jìn)程在當(dāng)前節(jié)點(diǎn)的編號,因?yàn)槲覀冊O(shè)置了2個進(jìn)程,因此進(jìn)程的編號就是0和1
在很多博客中都直接說明local_rank等于進(jìn)程內(nèi)的GPU編號,這種說法實(shí)際上是不準(zhǔn)確的。這個編號并不是GPU的編號??!
在使用啟動命令時,torch.distributed.launch工具會默認(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()
# 注意這個參數(shù),必須要以這種形式指定,即使代碼中不使用。因?yàn)?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])
可以看到,此時的local_rank共有4個,與進(jìn)程數(shù)相同。并且我們設(shè)置的可用GPU的id是1,2,3,4,而local_rank的輸出為0,1,2,3,可見local_rank并不是GPU的編號。
雖然在代碼中模型并行的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
此時我們再使用nvidia-smi來查看GPU的使用情況,如下??梢钥吹酱藭r使用的GPU就是local rank的id。相比于情形2,我們可以總結(jié):
當(dāng)沒有設(shè)置可用的GPU ID時,所采用的GPU id就等于local rank的id。本質(zhì)上是將進(jìn)程的編號作為GPU編號使用,因此local_rank等于進(jìn)程的編號這個定義是不變的。
當(dāng)設(shè)置可用的GPU ID,所采用的GPU id就等于GPU id。

情形4:將nproc_per_node設(shè)置為5,即超出了可以用的GPU數(shù)
輸出結(jié)果如下,可以看到是報(bào)錯的,因?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)練的兩種方式的文章就介紹到這了,更多相關(guān)pytorch的多GPU訓(xùn)練內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家!
相關(guān)文章
python多線程性能測試之快速mock數(shù)據(jù)
這篇文章主要為大家介紹了python多線程性能測試之快速mock數(shù)據(jù),有需要的朋友可以借鑒參考下,希望能夠有所幫助,祝大家多多進(jìn)步,早日升職加薪2022-06-06
Python實(shí)現(xiàn)獲取nginx服務(wù)器ip及流量統(tǒng)計(jì)信息功能示例
這篇文章主要介紹了Python實(shí)現(xiàn)獲取nginx服務(wù)器ip及流量統(tǒng)計(jì)信息功能,涉及Python針對nginx服務(wù)器信息操作相關(guān)實(shí)現(xiàn)技巧,需要的朋友可以參考下2018-05-05
python實(shí)現(xiàn)經(jīng)典排序算法的示例代碼
這篇文章主要介紹了python實(shí)現(xiàn)經(jīng)典排序算法的示例代碼,文中通過示例代碼介紹的非常詳細(xì),對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值,需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2021-02-02
Python實(shí)現(xiàn)快速將pdf文件剪切成多個圖片
這篇文章主要為大家詳細(xì)介紹了如何使用Python實(shí)現(xiàn)快速將pdf文件剪切成多個圖片,文中的示例代碼講解詳細(xì),有需要的小伙伴可以跟隨小編一起學(xué)習(xí)一下2024-01-01
Python增量循環(huán)刪除MySQL表數(shù)據(jù)的方法
這篇文章主要介紹了Python增量循環(huán)刪除MySQL表數(shù)據(jù)的相關(guān)資料,本文介紹的非常詳細(xì),具有參考借鑒價值,需要的朋友可以參考下2016-09-09

