實(shí)踐教程 | Pytorch的nn.DataParallel詳細(xì)解析

前言

pytorch中的GPU操作默認(rèn)是異步的，當(dāng)調(diào)用一個(gè)使用GPU的函數(shù)時(shí)，這些操作會(huì)在特定設(shè)備上排隊(duì)但不一定在稍后執(zhí)行。這就使得pytorch可以進(jìn)行并行計(jì)算。但是pytorch異步計(jì)算的效果對(duì)調(diào)用者是不可見的。

但平時(shí)我們用的更多其實(shí)是多GPU的并行計(jì)算，例如使用多個(gè)GPU訓(xùn)練同一個(gè)模型。Pytorch中的多GPU并行計(jì)算是數(shù)據(jù)級(jí)并行，相當(dāng)于開了多個(gè)進(jìn)程，每個(gè)進(jìn)程自己獨(dú)立運(yùn)行，然后再整合在一起。

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

注：多GPU計(jì)算的前提是你的計(jì)算機(jī)上得有多個(gè)GPU,在cmd上輸入nvidia-smi來查看自己的設(shè)備上的GPU信息。

nn.DataParallel詳細(xì)解析

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0):

這個(gè)函數(shù)主要有三個(gè)參數(shù)：

1. module：即模型，此處注意，雖然輸入數(shù)據(jù)被均分到不同gpu上，但每個(gè)gpu上都要拷貝一份模型。
2. device_ids：即參與訓(xùn)練的gpu列表，例如三塊卡， device_ids = [0，1，2]。
3. output_device：指定輸出gpu，一般省略。在省略的情況下，默認(rèn)為第一塊卡，即索引為0的卡。此處有一個(gè)問題，輸入計(jì)算是被幾塊卡均分的，但輸出loss的計(jì)算是由這一張卡獨(dú)自承擔(dān)的，這就造成這張卡所承受的計(jì)算量要大于其他參與訓(xùn)練的卡。

一般我們使用torch.nn.DataParallel()這個(gè)函數(shù)來進(jìn)行，接下來我將用一個(gè)例子來演示如何進(jìn)行多GPU計(jì)算：

net = torch.nn.Linear(100,1)
print(net)
print('---------------------')
net = torch.nn.DataParallel(net, device_ids=[0,3])
print(net)

輸出：

Linear(in_features=10, out_features=1, bias=True)
---------------------
DataParallel(
  (module): Linear(in_features=10, out_features=1, bias=True)
)

可以看到nn.DataParallel()包裹起來了。然后我們就可以使用這個(gè)net來進(jìn)行訓(xùn)練和預(yù)測了，它將自動(dòng)在第0塊GPU和第3塊GPU上進(jìn)行并行計(jì)算，然后自動(dòng)的把計(jì)算結(jié)果進(jìn)行了合并。

下面來具體講講nn.DataParallel中是怎么做的：

首先在前向過程中，你的輸入數(shù)據(jù)會(huì)被劃分成多個(gè)子部分（以下稱為副本）送到不同的device中進(jìn)行計(jì)算，而你的模型module是在每個(gè)device上進(jìn)行復(fù)制一份，也就是說，輸入的batch是會(huì)被平均分到每個(gè)device中去，但是你的模型module是要拷貝到每個(gè)devide中去的，每個(gè)模型module只需要處理每個(gè)副本即可，當(dāng)然你要保證你的batch size大于你的gpu個(gè)數(shù)。然后在反向傳播過程中，每個(gè)副本的梯度被累加到原始模塊中。概括來說就是：DataParallel會(huì)自動(dòng)幫我們將數(shù)據(jù)切分 load 到相應(yīng) GPU，將模型復(fù)制到相應(yīng) GPU，進(jìn)行正向傳播計(jì)算梯度并匯總。

注意還有一句話，官網(wǎng)中是這樣描述的：

The parallelized module must have its parameters and buffers on device_ids[0] before running this [DataParallel](https://link.zhihu.com/?target=https%3A//pytorch.org/docs/stable/nn.html%3Fhighlight%3Dtorch%2520nn%2520datapa%23torch.nn.DataParallel) module.

意思就是：在運(yùn)行此DataParallel模塊之前，并行化模塊必須在device_ids [0]上具有其參數(shù)和緩沖區(qū)。在執(zhí)行DataParallel之前，會(huì)首先把其模型的參數(shù)放在device_ids[0]上，一看好像也沒有什么毛病，其實(shí)有個(gè)小坑。我舉個(gè)例子，服務(wù)器是八卡的服務(wù)器，剛好前面序號(hào)是0的卡被別人占用著，于是你只能用其他的卡來，比如你用2和3號(hào)卡，如果你直接指定device_ids=[2, 3]的話會(huì)出現(xiàn)模型初始化錯(cuò)誤，類似于module沒有復(fù)制到在device_ids[0]上去。那么你需要在運(yùn)行train之前需要添加如下兩句話指定程序可見的devices，如下：

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

當(dāng)你添加這兩行代碼后，那么device_ids[0]默認(rèn)的就是第2號(hào)卡，你的模型也會(huì)初始化在第2號(hào)卡上了，而不會(huì)占用第0號(hào)卡了。這里簡單說一下設(shè)置上面兩行代碼后，那么對(duì)這個(gè)程序而言可見的只有2和3號(hào)卡，和其他的卡沒有關(guān)系，這是物理上的號(hào)卡，邏輯上來說其實(shí)是對(duì)應(yīng)0和1號(hào)卡，即device_ids[0]對(duì)應(yīng)的就是第2號(hào)卡，device_ids[1]對(duì)應(yīng)的就是第3號(hào)卡。

當(dāng)然你要保證上面這兩行代碼需要定義在下面這兩行代碼之前，一般放在train.py中import一些package之后：

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

那么在訓(xùn)練過程中，你的優(yōu)化器同樣可以使用nn.DataParallel，如下兩行代碼：

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

nn.DataParallel一些常見問題解析

1.多GPU計(jì)算減少了程序運(yùn)行的時(shí)間？

很多同學(xué)發(fā)現(xiàn)在進(jìn)行多GPU運(yùn)算時(shí)，程序花費(fèi)的時(shí)間反而更多了，這其實(shí)是因?yàn)槟愕腷atch_size太小了，因?yàn)閠orch.nn.DataParallel()這個(gè)函數(shù)是將每個(gè)batch的數(shù)據(jù)平均拆開分配到多個(gè)GPU上進(jìn)行計(jì)算，計(jì)算完再返回來合并。這導(dǎo)致GPU之間的開關(guān)和通訊過程占了大部分的時(shí)間開銷。

大家可以使用watch \-n 1 nvidia-smi這個(gè)命令來查看每1s各個(gè)GPU的運(yùn)行情況，如果發(fā)現(xiàn)每個(gè)GPU的占用率均低于50%，基本可以肯定你使用多GPU計(jì)算所花的時(shí)間要比單GPU計(jì)算花的時(shí)間更長了。

2. 如何保存和加載多GPU網(wǎng)絡(luò)？

如何來保存和加載多GPU網(wǎng)絡(luò)，它與普通網(wǎng)絡(luò)有一點(diǎn)細(xì)微的不同：

net = torch.nn.Linear(10,1)  # 先構(gòu)造一個(gè)網(wǎng)絡(luò)
net = torch.nn.DataParallel(net, device_ids=[0,3])  #包裹起來
torch.save(net.module.state_dict(), './networks/multiGPU.h5') #保存網(wǎng)絡(luò)

# 加載網(wǎng)絡(luò)
new_net = torch.nn.Linear(10,1)
new_net.load_state_dict(torch.load("./networks/multiGPU.h5"))

因?yàn)?code style="font-size: 14px;word-wrap: break-word;padding: 2px 4px;border-radius: 4px;margin: 0 2px;color: #1e6bb8;background-color: rgba(27,31,35,.05);font-family: Operator Mono, Consolas, Monaco, Menlo, monospace;word-break: break-all;">DataParallel實(shí)際上是一個(gè)nn.Module，所以我們在保存時(shí)需要多調(diào)用了一個(gè)net.module，模型和優(yōu)化器都需要使用net.module來得到實(shí)際的模型和優(yōu)化器。

3. 為什么第一塊卡的顯存會(huì)占用的更多一些？？？

最后一個(gè)參數(shù)output_device一般情況下是省略不寫的，那么默認(rèn)就是在device_ids[0]，也就是第一塊卡上，也就解釋了為什么第一塊卡的顯存會(huì)占用的比其他卡要更多一些。

進(jìn)一步說也就是當(dāng)你調(diào)用nn.DataParallel的時(shí)候，只是在你的input數(shù)據(jù)是并行的，但是你的output loss卻不是這樣的，每次都會(huì)在第一塊GPU相加計(jì)算，這就造成了第一塊GPU的負(fù)載遠(yuǎn)遠(yuǎn)大于剩余其他的顯卡。

4. 直接使用nn.DataParallel的時(shí)候，訓(xùn)練采用多卡訓(xùn)練，會(huì)出現(xiàn)一個(gè)warning？？？

UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; 
will instead unsqueeze and return a vector.

首先說明一下：

每張卡上的loss都是要匯總到第0張卡上求梯度，更新好以后把權(quán)重分發(fā)到其余卡。但是為什么會(huì)出現(xiàn)這個(gè)warning，這其實(shí)和nn.DataParallel中最后一個(gè)參數(shù)dim有關(guān)，其表示tensors被分散的維度，默認(rèn)是0，nn.DataParallel將在dim0（批處理維度）中對(duì)數(shù)據(jù)進(jìn)行分塊，并將每個(gè)分塊發(fā)送到相應(yīng)的設(shè)備。單卡的沒有這個(gè)warning，多卡的時(shí)候采用nn.DataParallel訓(xùn)練會(huì)出現(xiàn)這個(gè)warning，由于計(jì)算loss的時(shí)候是分別在多卡計(jì)算的，那么返回的也就是多個(gè)loss，你使用了多少個(gè)gpu，就會(huì)返回多少個(gè)loss。（有人建議DataParallel類應(yīng)該有reduce和size_average參數(shù)，比如用于聚合輸出的不同loss函數(shù)，最終返回一個(gè)向量，有多少個(gè)gpu，返回的向量就有幾維。）

關(guān)于這個(gè)問題在pytorch官網(wǎng)的issues上有過討論，下面簡單摘出一些:

https://github.com/pytorch/pytorch/issues/9811github.com

前期探討中，有人提出求loss平均的方式會(huì)在不同數(shù)量的gpu上訓(xùn)練會(huì)以微妙的方式影響結(jié)果。模塊返回該batch中所有損失的平均值，如果在4個(gè)gpu上運(yùn)行，將返回4個(gè)平均值的向量。然后取這個(gè)向量的平均值。但是，如果在3個(gè)GPU或單個(gè)GPU上運(yùn)行，這將不是同一個(gè)數(shù)字，因?yàn)槊總€(gè)GPU處理的batch size不同！舉個(gè)簡單的例子（就直接摘原文出來）：

A batch of 3 would be calculated on a single GPU and results would be [0.3, 0.2, 0.8] and model that returns the loss would return 0.43.  

If cast to DataParallel, and calculated on 2 GPUs, [GPU1 - batch 0,1], [GPU2 - batch 2] - return values would be [0.25, 0.8] (0.25 is average between 0.2 and 0.3)- taking the average loss of [0.25, 0.8] is now 0.525!  

Calculating on 3 GPUs, one gets [0.3, 0.2, 0.8] as results and average is back to 0.43!

似乎一看，這么求平均loss確實(shí)有不合理的地方。那么有什么好的解決辦法呢，可以使用size_average=False，reduce=True作為參數(shù)。每個(gè)GPU上的損失將相加，但不除以GPU上的批大小。然后將所有平行損耗相加，除以整批的大小，那么不管幾塊GPU最終得到的平均loss都是一樣的。

那pytorch貢獻(xiàn)者也實(shí)現(xiàn)了這個(gè)loss求平均的功能，即通過gather的方式來求loss平均：

https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3begithub.com

如果它們在一個(gè)有2個(gè)GPU的系統(tǒng)上運(yùn)行，DP將采用多GPU路徑，調(diào)用gather并返回一個(gè)向量。如果運(yùn)行時(shí)有1個(gè)GPU可見，DP將采用順序路徑，完全忽略gather，因?yàn)檫@是不必要的，并返回一個(gè)標(biāo)量。

參考鏈接

1. Pytorch多GPU計(jì)算之torch.nn.DataParallel：https://blog.csdn.net/wangkaidehao/article/details/104411682
2. Pytorch的nn.DataParallel：https://zhuanlan.zhihu.com/p/102697821

如果覺得有用，就請(qǐng)分享到朋友圈吧！