百度飞桨大规模分类库PLSC—单机训练6000万类视觉分类模型

　　今天给大家分享的是百度飞桨大规模分类库PLS实现了，单机训练6000万类视觉分类模型，下面我们来看具体详情！

　　以下文章来源于飞桨PaddlePaddle，作者飞桨PaddlePaddle

　　大规模分类任务

　　在介绍大规模分类任务之前，我们先简短回顾一下通常的分类任务。大家熟知的视觉分类任务中，网络模型由特征提取器（Backbone）和分类器（Classifier）组成。分类的类别数有2类（如，前景/背景分类）、10类（如，MNIST数据分类）、80类（如，COCO数据分类）和1000类（如，ImageNet数据分类）等等。比较主流的特征提取器有ResNet、MobileNet等网络结构，分类器则通常采用线性分类层（全连接层，FC）。

　　大规模分类任务的『大规模』指模型参数规模非常大，包括以下3种情况：特征提取器参数规模大、分类器参数规模大，以及两者参数规模都大。万物互联时代，随着人工智能、5G和IoT等技术的发展，分类模型分类的类别数不断增加，类别数可以达到上千万甚至更多。在这种背景下，分类网络模型的FC层的类别数增加，参数规模爆炸式增长。基于度量学习的分类模型，通常在训练阶段使用闭集数据集学习特征提取器和分类器，在推理阶段，仅使用特征提取器提取输入图像的特征，并与预提取的特征进行相似性对比得出是否属于同一类，如图1所示。

　　大规模分类网络模型训练和部署示意图

　　大规模分类模型训练难点

　　本文聚焦于解决大规模分类模型训练问题。有小伙伴可能会问：大规模分类不就是一个普通的图像分类吗，除了分类类别数较多导致的FC层参数量大以外，还有什么难题？图像分类领域每年有大量的论文和工作在ImageNet数据集上取得新的SOTA，随便从Github上找个图像分类库来训练是不是就可以了？

　　然而，FC层参数规模的急剧增长在训练时会带来以下两方面挑战：

　　首先是存储问题。假设分类类别数为4000万，在训练阶段特征向量的维度为512，并且以32比特浮点数存储模型参数，那么仅FC层的参数就可达512*40000000*4（bytes）/1024/1024=76.29GB，远远超出主流显卡的存储容量。这是普通图像分类库无法解决的存储问题。

　　其次是速度问题。普通分类模型也面临同样的问题。随着训练数据、模型规模和分类类别数的增加，模型训练的复杂度显著增长，导致模型训练所需要的时间不断增长。速度是人类永无止境的追求，如何在更短的时间内训练大规模分类模型也是工程实践中迫切需要解决的问题。

　　为了解决以上两方面难题，学术界和工业界不断围绕着训练的显存消耗和速度进行优化；飞桨团队也持续不断地打磨升级大规模分类库PLSC（Paddle Large Scale Classification），提供数据并行&模型并行混合训练、类别中心采样、稀疏梯度参数更新和FP16训练等解决方案。

　　解决方案详解

　　接下来，小编将和大家一起来分享PLSC提供的数据并行&模型并行混合训练、模型并行Loss计算、类别中心采样、稀疏梯度参数更新和FP16训练解决方案。

　　混合并行训练

　　为了提高训练效率，通常使用多张GPU卡做数据并行训练。但是对于大规模分类问题，类别数非常多，导致单卡无法训练。例如，4000万类的分类网络模型仅FC层的参数量就高达76.29GB，远远超出主流显卡的存储容量。

　　此外，数据并行下FC层的梯度通信量也巨大，导致训练速度难以接受。针对FC层参数存储和梯度通信问题，我们自然会想到是否可以将参数存放到多张GPU卡上？答案是肯定的。我们可以采用模型并行策略，将FC层参数切分到多张卡上。如图2所示，Backbone网络采用数据并行，分类FC层采用模型并行，兼顾了数据并行的训练效率和分类FC层参数的存储和梯度通信需求。在分类类别数为4000万类时，假如使用单机8卡，那么每张卡上的FC层仅需存放500万类，参数的存储大小为76.29GB/8=9.54GB。

　　采用数据并行和模型并行结合的方式，单机8卡情形下前向计算过程如下：

　　1.每张卡接受一个Batch的数据，假设batch size为64；

　　2.每张卡使用输入数据做数据并行计算，经过Backbone得到512维特征向量，维度为64x512；

　　3.对每张卡上的特征和标签做allgather操作，从其他卡上收集特征和标签，此时每张卡上拥有全量特征维度为512x512，全量标签维度为512x1；

　　4.全量特征（512x512）和部分FC参数（512x5000000）做矩阵乘操作得到logits，维度为512x5000000；

　　5.使用模型并行的SoftmaxWithCrossEntropy Loss函数计算loss。

　　Backbone数据并行&Classifer模型并行

　　模型并行Loss计算——API级MarginLoss函数

　　在度量学习领域中，ArcFace[2]论文将ArcFace，CosFace[3]和SphereFace[4]Loss函数用如下统一的公式表示，我们称为MarginLoss：

　　MarinLoss函数是在logits上增加了margin，最终基于的仍然是SoftmaxWithCrossEntropy Loss函数。模型并行下最容易想到的计算Loss的方法是用通信操作从其他卡上获取全量的logits。但是这种方法不仅需要巨大的通信量，同时需要临时存储其他卡上的logits，带来巨大的显存开销。

　　飞桨框架在模型并行策略下提供了对MarginLoss--paddle.nn.functional.margin_cross_entropy的原生支持。该接口通信量少且显存开销较小。图3给出模型并行下SoftmaxwithCrossEntropy的计算过程。首先，在每张卡上逐行计算logits的最大值，然后通过allreduce操作获取全局每行最大值。为了保持数值计算的稳定性，每行减去其对应的全局最大值。接着，逐行计算分母和，然后通过allreduce操作获取全局和值。最后，逐行计算loss,并通过allreduce操作获取全局loss。图中，我们对Loss和Softmax Probability计算做了共同表达式提取。

　　模型并行SoftmaxwithCrossEntropy计算过程

　　类别中心采样——API级支持PartialFC

　　采用数据并行&模型并行解决了FC参数存储的问题。但是，可以发现前向计算的logits存储需求也非常大，在混合并行训练小节的假设条件下为512*5000000*4（bytes）/1024/1024=9.54 GB。考虑前向计算、反向计算和参数更新相关变量，当优化方法使用Momentum时，可以得到FC层需要的存储大小：

　　其中d表示特征的维度，c表示总类别数，k表示GPU卡数，N表示Batch大小，Memw表示参数存储大小，Memlogits表示logits存储大小，MemFc表示FC层总的存储大小。当类别数增大时，我们可以将FC层参数切分到不同卡上，以保持每张卡上存储的参数大小不变。然而，logits的维度却是随卡数线性增长的。

　　因此，卡数增大k倍，Memlogits也增大k倍。训练过程中，FC层总的存储大小MemFc等于3倍Memw（weight，gradient和velocity）加2倍Memlogits（activation和gradient）。为了解决logits和对应的梯度存储问题，PartialFC[5]提出基于采样的FC层，即从全量FC中采样一部分类别中心参与本次迭代的学习。如图4所示，对于正样本对应的类别中心全部保留，对于负类别中心按比例随机采样。

　　假设采样比例为1/10，则logits的维度为512x500000，存储大小为0.1*9.54 GB=0.954GB。优化前存储大小为2*9.54 GB=19.08 GB，采用PartialFC时需要的显存开销为2*0.954GB=1.908GB，可见使用PartialFC可以大幅减小显存开销。

　　PatialFC采样过程

　　飞桨提供了原生支持上述采样过程的API——

　　paddle.nn.functional.class_center_sample。

　　稀疏梯度参数更新——SparseMomentum

　　稀疏梯度参数更新是PLSC一大亮点。虽然PartialFC通过采样的方法缓解了logits和对应梯度的显存开销，但是通过上述分析我们发现FC层仍然需要3倍的Mem_w，分别对应FC层参数、梯度和优化器状态。我们是否可以进一步优化显存呢？答案是肯定的。

　　如图5左所示，在前向计算过程中，通过对参数W的采样，得到采样类别中心Wsub；反向计算梯度时，首先得到稀疏梯度Wsub grad，进而得到参数梯度W grad；在参数更新阶段，momentum算子使用传入的参数Wsub，参数梯度W grad和优化器状态W velocity更新参数W和优化器状态W velocity。我们通过分析发现，参数梯度W grad是冗余的，其可以通过稀疏梯度Wsub grad得到。为此，我们设计和开发了sparse_momentum接口。相比于momentum，该接口需要额外传入参数index，表示FC参数采样的索引值，计算过程如图5右所示。使用该接口，可以大幅减少梯度的存储空间，从而可以训练更大规模参数的模型。相比momentum需要3*9.54Gb+2*0.954Gb=30.528GB的存储空间，使用sparse_momentum，FC层仅需要2*9.54Gb+2*0.954Gb=20.988GB的存储空间，显存空间降低31.25%。

　　Momentum和SparseMomentum的更新过程对比

　　FP16训练——节省50%显存

　　PLSC的另外一个亮点是使用FP16训练，即整个训练过程中参数、Activation、梯度和优化器状态均使用FP16，相比于FP32显存空间节省50%。相比FP32和AMP[6]，FP16可以大幅节省显存，并大幅提升训练速度。

　　图6分别给出是FP32，AMP和FP16的计算过程。FP32计算过程中，所有模型参数、梯度、Activation和优化器状态的数据类型均为FP32。AMP计算过程中，模型参数和优化器状态为FP32；计算过程中，将参数cast成FP16，因此Activation和梯度也是FP16；优化阶段参数梯度需要重新cast为FP32；所以，相比于FP32，AMP通过将Activation和对应的梯度存储为FP16，节省显存开销。PLSC使用的是真正意义的FP16，即模型参数、Activation、梯度和优化器状态均使用FP16；相比FP32显存开销减少50%；此外，由于消除了cast操作，FP16相比于AMP可以进一步提升训练速度。

　　FP32、AMP和FP16计算过程

　　实验结果

　　在上一节我们介绍了大规模分类模型训练的一些解决方案，这些解决方案都已经在PLSC中实现且开源。此外，我们也已经将PLSC开源到人脸识别社区InsighFace[1]。

　　PLSC库地址：https://github.com/PaddlePaddle/PLSC

　　PLSC具有以下亮点：

　　高吞吐，低显存，简单易用；

　　支持单机和多机分布式训练，API级支持模型并行、PartialFC和MarginLoss；

　　支持FP16训练；

　　同时支持静态图和动态图。

　　接下来我们将从训练精度、显存开销和训练速度3个维度来评测PLSC。

　　MS1MV3精度

　　下表给出MS1MV3数据上的精度对比。

　　表1不同框架实现的Repo在MS1MV3数据集上的精度对比

　　从表1中，我们可以看到虽然PLSC使用了FP16，但在主要数据集上PLSC的精度仍然可以打平其它框架实现。

　　最大支持类别数

　　实验配置：

　　GPUs：8 NVIDIA Tesla V100 32G；

　　BatchSize：64/512（每张卡的batch size是64，全局batch size是512个样本）；

　　SampleRatio：0.1（PartialFC采样率为0.1）。

　　表2不同框架实现支持的最大类别数对比

　　表中数据说明，相比其他的框架实现，PLSC在显存优化方面具有显著优势：静态图最多支持6000万类别分类，动态图最多可支持6700万类别分类。

　　吞吐量对比

　　吞吐量是每秒训练的样本数。我们采用公开数据集MS1MV3来测试。为了取得稳定和公平的结果，我们对每个实验配置都运行了5组实验，每组运行200个steps，然后计算50到150个steps间的100个steps吞吐量的平均值，最后得到这5组实验的平均吞吐量后再取中位数作为最终的吞吐量结果。以下是实验配置：

　　Tesla V100(32G)：Driver Version：450.80.02,CUDA Version：11.0；

　　Tesla A100(40G)：Driver Version：460.32.03,CUDA Version：11.2；

　　Datasets：MS1MV3（93431l类）；

　　SampleRatio：0.1（使用了PartialFC，采样率为0.1）；

　　BatchSize：128（每张卡128个样本）。

　　图7：不同框架实现的Repo吞吐量对比

　　从图7可以看出PLSC静态图模式下，优于所有框架实现，尤其在A100，ResNet50，FP16，8卡的配置下，PLSC吞吐量高达9500imgs/s。

　　直播课深入解读

　　大规模分类的一个典型应用就是人脸识别，这里我们邀请到了InsightFace项目发起人与飞桨开发者一起，在B站直播间，为大家深入技术解读。

　　项目地址：

　　GitHub：https://github.com/PaddlePaddle/PLSC

　　GitHub:https://github.com/deepinsight/insightface

　　更多百度飞桨PLS相关内容，百度云服务中心持续分享中！

　　推荐阅读：百度语音合成长文本在线合成上线邀测啦