千万级类别人脸识别模型并行训练
作者:互联网
并行训练的方式:
- 1.nn.DataParallel数据并行。将一个batchsize中的数据分给多个GPU并行训练。
- 2.模型并行。将FC层拆分给多个GPU进行并行训练。
- 3.partial_fc。(抽样fc层)
一、模型并行
目前处理大规模(数据多、类别大)数据集的方法:
混合并行:即backbone使用数据并行,分类层使用模型并行;
该方法具备两个优点:
- 1)缓解了 W 的存储压力。将W划分为k个子矩阵w;
- 2)将 W 梯度的通信转换成了所有GPU的特征 X 与 softmax 局部分母的通信,大大降低了数据并行带来的通信开销。
模型并行的结构图:
模型并行方法的弊端:
模型并行的方式理论上看似能无限增加类别数(只要增加GPU数量即可),但是实际上大家在尝试更大规模、更多机器的时候,会发现显存不够用了,好像增加类别数的同时增加机器,单个GPU的显存还在增长?其实我们忽略了另外一个占据显存的张量:predicted logits的存储会受到总批大小的增加的影响。logits(预计日志的存储会受到总批大小的增加的影响)。
首先定义 logits = X_w,其中 w 为存储在每张GPU上的子矩阵,X 为经过集合通信 Allgaher 收集到的全局特征,d 为特征的维度大小,C 为总的类别数,k 为GPU的个数。其中每块GPU中 w 占用的显存为:
结论:当我们不停的增加GPU数量时,logits占用的内存也会增大,当GPU数量K大到一定量时,内存就会溢出了。
二、Partial FC(FC抽样)
该研究对此提出了一个简单的解决方案:
在实现混合并行时,不仅同步每张卡的特征,同时也同步每张卡的标签,这样每张卡都具备所有卡的完整特征和标签。假设总的批次大小为 kN,则至多会有 kN 个正类中心随机分布在所有的GPU中,让每个正类中心所属的GPU将该正类采样出来即可,每张GPU正类采出来后,再随机用负类补齐到约定的采样率。这样每张GPU采样到的类中心一样多,实现负载均衡。后续的过程就是分类层的模型并行部分了,需要注意的是,只有采样出来的类中心的权重和动量会更新。
partial_fc论文博客:
https://jishuin.proginn.com/p/763bfbd2fee4
https://blog.csdn.net/weixin_43152063/article/details/115307938
https://blog.51cto.com/u_15282017/2974039
https://blog.csdn.net/zengwubbb/article/details/109050165
标签:显存,人脸识别,模型,并行,千万级,每张,正类,GPU 来源: https://blog.csdn.net/chen1234520nnn/article/details/119008997