Pytorch——多卡GUP训练原理(torch.nn.DataParallel)
本文将记录如何使用单机多卡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博客
