單機多卡訓練,即并行訓練。并行訓練又分為數據并行 (Data Parallelism) 和模型并行兩種。
數據并行指的是,多張 GPU 使用相同的模型副本,但是使用不同的數據批進行訓練。而模型并行指的是,多張GPU 分別訓練模型的不同部分,使用同一批數據。
數據并行
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
【參數】
module : 要進行并行的 module。這里隱含了一點 ,即網絡中的某一層也是可以進行數據并行的,但是一般不會這么使用。
device_ids : CUDA 列表,可以為 torch.device 類型,也可以是編號組成的 int 列表。默認使用全部 GPU
output_device : 某一 GPU 編號或 torch.device 。指定輸出的 GPU,默認為第一個,即 device_ids[0]
數據并行的具體原理流程為:
1.將模型加載至主設備上,作為 controller,一般設置為 cuda:0
在每次迭代時,執行如下操作:
2.將 controller 模型復制(broadcast)到每一個指定的 GPU 上
將總輸入的數據 batch,進行均分,分別作為各對應副本的輸入 (scatter)
3.每個副本獨立進行前向傳播,并進行反向傳播,但只是求取梯度
4.將各副本的梯度匯總(gather)到 controller 設備,并進行求和 (reduced add) During the backwards pass, gradients from each replica are summed into the original module.
5.更具總體度,更新 controller 設備上的參數