首先查看一下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的话，只能采用模型并行的方法进行训练。
将模型分成不同的部分，存在不同的显卡上。前一个的输出是下一个的输入。
模型并行
