单机多卡训练,即并行训练。并行训练又分为数据并行 (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 设备上的参数