公司配備多卡的GPU服務(wù)器，當(dāng)我們在上面跑程序的時候，當(dāng)?shù)螖?shù)或者epoch足夠大的時候，我們通常會使用nn.DataParallel函數(shù)來用多個GPU來加速訓(xùn)練。一般我們會在代碼中加入以下這句：device_ids = [0, 1]net = torch.nn.DataParallel(net, device_ids=device_ids)
似乎只要加上這一行代碼，你在ternimal下執(zhí)行watch -n 1 nvidia-smi后會發(fā)現(xiàn)確實會使用多個GPU來并行訓(xùn)練。但是細心點會發(fā)現(xiàn)其實第一塊卡的顯存會占用的更多一些，那么這是什么原因?qū)е碌模坎殚唒ytorch官網(wǎng)的nn.DataParrallel相關(guān)資料，首先我們來看下其定義如下：CLASS torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
其中包含三個主要的參數(shù)：module，device_ids和output_device。官方的解釋如下：module即表示你定義的模型，device_ids表示你訓(xùn)練的device，output_device這個參數(shù)表示輸出結(jié)果的device，而這最后一個參數(shù)output_device一般情況下是省略不寫的，那么默認就是在device_ids[0]，也就是第一塊卡上，也就解釋了為什么第一塊卡的顯存會占用的比其他卡要更多一些。進一步說也就是當(dāng)你調(diào)用nn.DataParallel的時候，只是在你的input數(shù)據(jù)是并行的，但是你的output loss卻不是這樣的，每次都會在第一塊GPU相加計算，這就造成了第一塊GPU的負載遠遠大于剩余其他的顯卡。
下面來具體講講nn.DataParallel中是怎么做的。首先在前向過程中，你的輸入數(shù)據(jù)會被劃分成多個子部分（以下稱為副本）送到不同的device中進行計算，而你的模型module是在每個device上進行復(fù)制一份，也就是說，輸入的batch是會被平均分到每個device中去，但是你的模型module是要拷貝到每個devide中去的，每個模型module只需要處理每個副本即可，當(dāng)然你要保證你的batch size大于你的gpu個數(shù)。然后在反向傳播過程中，每個副本的梯度被累加到原始模塊中。概括來說就是：DataParallel 會自動幫我們將數(shù)據(jù)切分 load 到相應(yīng) GPU，將模型復(fù)制到相應(yīng) GPU，進行正向傳播計算梯度并匯總。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]上，一看好像也沒有什么毛病，其實有個小坑。我舉個例子，服務(wù)器是八卡的服務(wù)器，剛好前面序號是0的卡被別人占用著，于是你只能用其他的卡來，比如你用2和3號卡，如果你直接指定device_ids=[2, 3]的話會出現(xiàn)模型初始化錯誤，類似于module沒有復(fù)制到在device_ids[0]上去。那么你需要在運行train之前需要添加如下兩句話指定程序可見的devices，如下：os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"
當(dāng)你添加這兩行代碼后，那么device_ids[0]默認的就是第2號卡，你的模型也會初始化在第2號卡上了，而不會占用第0號卡了。這里簡單說一下設(shè)置上面兩行代碼后，那么對這個程序而言可見的只有2和3號卡，和其他的卡沒有關(guān)系，這是物理上的號卡，邏輯上來說其實是對應(yīng)0和1號卡，即device_ids[0]對應(yīng)的就是第2號卡，device_ids[1]對應(yīng)的就是第3號卡。當(dāng)然，你要保證上面這兩行代碼需要定義在device_ids = [0, 1]net = torch.nn.DataParallel(net, device_ids=device_ids)
這兩行代碼之前，一般放在train.py中import一些package之后。那么在訓(xùn)練過程中，你的優(yōu)化器同樣可以使用nn.DataParallel，如下兩行代碼：optimizer = torch.optim.SGD(net.parameters(), lr=lr)optimizer = nn.DataParallel(optimizer, device_ids=device_ids)
那么使用nn.DataParallel后，事實上DataParallel也是一個Pytorch的nn.Module，那么你的模型和優(yōu)化器都需要使用.module來得到實際的模型和優(yōu)化器，如下：保存模型：torch.save(net.module.state_dict(), path)加載模型：net=nn.DataParallel(Resnet18())net.load_state_dict(torch.load(path))net=net.module優(yōu)化器使用：optimizer.step() --> optimizer.module.step()
還有一個問題就是，如果直接使用nn.DataParallel的時候，訓(xùn)練采用多卡訓(xùn)練，會出現(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張卡上求梯度，更新好以后把權(quán)重分發(fā)到其余卡。但是為什么會出現(xiàn)這個warning，這其實和nn.DataParallel中最后一個參數(shù)dim有關(guān)，其表示tensors被分散的維度，默認是0，nn.DataParallel將在dim0（批處理維度）中對數(shù)據(jù)進行分塊，并將每個分塊發(fā)送到相應(yīng)的設(shè)備。單卡的沒有這個warning，多卡的時候采用nn.DataParallel訓(xùn)練會出現(xiàn)這個warning，由于計算loss的時候是分別在多卡計算的，那么返回的也就是多個loss，你使用了多少個gpu，就會返回多少個loss。（有人建議DataParallel類應(yīng)該有reduce和size_average參數(shù)，比如用于聚合輸出的不同loss函數(shù)，最終返回一個向量，有多少個gpu，返回的向量就有幾維。）關(guān)于這個問題在pytorch官網(wǎng)的issues上有過討論，下面簡單摘出一些。鏈接：https://github.com/pytorch/pytorch/issues/9811前期探討中，有人提出求loss平均的方式會在不同數(shù)量的gpu上訓(xùn)練會以微妙的方式影響結(jié)果。模塊返回該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/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be如果它們在一個有2個GPU的系統(tǒng)上運行，DP將采用多GPU路徑，調(diào)用gather并返回一個向量。如果運行時有1個GPU可見，DP將采用順序路徑，完全忽略gather，因為這是不必要的，并返回一個標(biāo)量。其實關(guān)于多卡訓(xùn)練還有DistributedDataParallel等其他方式，就等實驗后繼續(xù)補充啦。
