Repository files navigation

English version

在人脸和re-id领域，部分私有的数据集的label数量可达上百万/千万/亿的规模，此时fc层的参数量就足以把显卡容量撑满，导致只能使用较小的batch_size，训练速度较慢，效果不佳

classFullyConnected(nn.Module):def__init__(self,in_dim,out_dim,num_gpu,model_parallel=False):super(FullyConnected,self).__init__()self.num_gpu=num_gpuself.model_parallel=model_parallelifmodel_parallel:self.fc_chunks=nn.ModuleList()foriinrange(num_gpu):_class_num=out_dim//num_gpuifi< (out_dim%num_gpu):_class_num+=1self.fc_chunks.append(nn.Linear(in_dim,_class_num,bias=False).cuda(i)                )else:self.classifier=nn.Linear(in_dim,out_dim,bias=False)defforward(self,x):ifself.model_parallel:x_list= []foriinrange(self.num_gpu):_x=self.fc_chunks[i](x.cuda(i))x_list.append(_x)x=torch.cat(x_list,dim=1)returnxelse:returnself.classifier(x)

类似的实现在这个基于PyTorch的人脸项目也能够看到

这个方案能够部分解决这个问题，但是又会引入新的问题：显存占用不均衡。由于最后结果依然要concat回0号卡上，且loss计算依然在0号卡上，0号卡的显存占用以及计算负载显著高于其他卡。受制于此，依然无法使用较大的batch_size

不仅解决了，还扩展到了更多场景下，支持人脸和re-id训练中常见的margin loss，支持混合精度训练与分布式训练。

几点小小的优势：

• 显存与计算负载合理分担到每张卡上，能够使用非常大的batch_size，训练得更加开心
• 只需做一些小小的修改就可以适配主流的margin loss，如ArcFace、SphereFace、CosFace、AM-softmax等等
• 相同的setting下对训练精度无影响（有数学推导保证结果正确）
• 在某些情况下甚至能加速训练（得益于优化后CrossEntropyLoss计算的过程中通信开销的降低）
• 支持混合精度训练与分布式训练

首先确认下你是否有必要使用模型并行：

• 数据集label规模是否在百万级以上？
• 模型的最后一层是否为fc层，是否使用CrossEntropyLoss？
• 显卡数量是否足够？（至少4~8张显卡）

如果以上答案均为肯定，那么你可以考虑使用模型并行。但是由于模型并行需要hack model和optimizer（分布式条件下更为复杂），目前需要自行移植到你的项目中。

• 普通的及混合精度训练可参考master分支
• 分布式训练可参考dist分支，目前仍在完善中

原理都是相通的，其他框架如MXNet甚至有对分布式支持更为友好的kvstore可供使用

• http://bindog.github.io/blog/2019/09/05/gpu-memory-balanced-model-parallel/

• http://bindog.github.io/blog/2020/04/12/model-parallel-with-apex/