本文將記錄如何使用單機多卡GPU的方式進行訓練，主要是采用DP模式（DDP模式一般用於多機多卡訓練）。
1、DP模式基本原理
DP模型進行單機多卡訓練基本步驟如下：
（1）將模型復制到各個GPU中，並將一個batch的數據划分成mini_batch並分發給給個GPU
（2）各個GPU獨自完成mini_batch的前向傳播，並把獲得的output傳遞給GPU_0（主GPU）
（3）GPU_0整合各個GPU傳遞過來的output，並計算loss。此時GPU_0可以對這些個loss進行一些聚合操作。
（4）GPU_0歸並loss之后，並進行后向傳播以及梯度下降從而完成模型參數的更新（此時只有GPU_0上的模型參數得到了更新），GPU_0將更新好的模型參數又傳遞給其余GPU
以上就是DP模式下多卡GPU進行訓練的方式。其實可以看到GPU_0不僅承擔了前向傳播的任務，還承擔了收集loss，並進行梯度下降。因此在使用DP模式進行單機多卡GPU訓練的時候會有一張卡的顯存利用會比其他卡更多，那就是你設置的GPU_0。
2、Pytorch如何進行單機多卡訓練
Pytorch中可以直接使用torch.nn.DataParallel(model)就可以完成了。具體的代碼模塊見下圖。
import torch.nn as nn
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '2,3,1' # 指定該程序可以識別的物理GPU編號
model = nn.DataParallel(model) # 在執行該語句之前最好加上model.cuda(),保證你的模型存在GPU上即可
第三行代碼建議一定要加上，這一行代碼是為了指定有哪些物理GPU是對該程序可見的，並且此處的順序將會是該程序在執行的時候識別GPU的順序（例如此處該程序的GPU_0=物理GPU2號，GPU_1=物理GPU3號，GPU_2=物理GPU1號），所以物理GPU2號將會作為該程序的主要GPU（用於計算loss，梯度下降，參數更新）。
對於數據，我們只需要按照平常的方式使用.cuda（）放置在GPU上即可，內部batch的拆分已經被封裝在了DataParallel模塊中。要注意的是，由於我們的model被nn.DataParallel()包裹住了，所以如果想要儲存模型的參數，需要使用model.module.state_dict()的方式才能取出（不能直接是model.state_dict()）。
參考網站：
pytorch多gpu並行訓練 - 知乎 (zhihu.com)
解決了PyTorch 使用torch.nn.DataParallel 進行多GPU訓練的一個BUG:模型(參數)和數據不在相同設備上_senius的博客-CSDN博客
pytorch遇到的bug記錄—2_羊藤枝的博客-CSDN博客
【分布式訓練】單機多卡的正確打開方式（三）：PyTorch - 知乎 (zhihu.com)
分布式訓練 - 單機多卡（DP和DDP）_love1005lin的博客-CSDN博客