pytorch 多GPU训练

首先查看一下GPU的数量

nvidia-smi

1 单机多卡训练

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
指定编号为0,1,2的GPU

模型加载到GPU上。

model = torch.nn.DataParallel(model)
model = model.cuda()

数据加载到GPU上

inputs = inputs.cuda()
labels = labels.cuda()

查看GPU的数量

torch.cuda.device_count()

1.1 解决多GPU训练显存使用不均衡的问题

采用DistributedDataParallel.

2 基本过程

-首先把模型加载到一个主设备
-把模型只读复制到多个设备
-把大的batch 数据也等分到不同的设备
-最后将所有设备计算得到的梯度合并更新主设备上的模型参数

多卡的batch大小=单卡batch*显卡数量
数据和初始网络都要在多个GPU中的第一块上就行。

每张卡独立forward自己那份data，但是在backward的时候，每张卡的额梯度会汇总到原始的module上，以此来实现并行训练。
但是这样的方法会造成原始Module所在的那张卡的显存压力比其他的要大，也就是存在负载不均衡的情况。

数据并行

3 模型并行

当单张显卡连一个样本都无法feed的话，只能采用模型并行的方法进行训练。
将模型分成不同的部分，存在不同的显卡上。前一个的输出是下一个的输入。

模型并行

参考