實踐教程 | Pytorch的nn.DataParallel詳細解析

前言

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

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

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

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

nn.DataParallel詳細解析

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

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

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

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

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()包裹起來了。然后我們就可以使用這個net來進行訓練和預測了，它將自動在第0塊GPU和第3塊GPU上進行并行計算，然后自動的把計算結果進行了合并。

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

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

注意還有一句話，官網(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.

意思就是：在運行此DataParallel模塊之前，并行化模塊必須在device_ids [0]上具有其參數(shù)和緩沖區(qū)。在執(zhí)行DataParallel之前，會首先把其模型的參數(shù)放在device_ids[0]上，一看好像也沒有什么毛病，其實有個小坑。我舉個例子，服務器是八卡的服務器，剛好前面序號是0的卡被別人占用著，于是你只能用其他的卡來，比如你用2和3號卡，如果你直接指定device_ids=[2, 3]的話會出現(xiàn)模型初始化錯誤，類似于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號卡。

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

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

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

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

nn.DataParallel一些常見問題解析

1.多GPU計算減少了程序運行的時間？

很多同學發(fā)現(xiàn)在進行多GPU運算時，程序花費的時間反而更多了，這其實是因為你的batch_size太小了，因為torch.nn.DataParallel()這個函數(shù)是將每個batch的數(shù)據(jù)平均拆開分配到多個GPU上進行計算，計算完再返回來合并。這導致GPU之間的開關和通訊過程占了大部分的時間開銷。

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

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

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

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

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

因為DataParallel實際上是一個nn.Module，所以我們在保存時需要多調(diào)用了一個net.module，模型和優(yōu)化器都需要使用net.module來得到實際的模型和優(yōu)化器。

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

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

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

4. 直接使用nn.DataParallel的時候，訓練采用多卡訓練，會出現(xiàn)一個warning？？？

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

首先說明一下：

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

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

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

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

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

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

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

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

參考鏈接

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

如果覺得有用，就請分享到朋友圈吧！