本文介绍了如何使用OneFlow框架实现大规模人脸识别，对比了与MXNet的实现，揭示了OneFlow在数据并行、模型并行以及Partial FC采样上的优势。OneFlow通过数据并行和模型并行的混合并行方式，解决了显存爆炸问题，支持超大规模人脸ID。在性能上，OneFlow在数据并行时速度是MXNet的2.82倍，模型并行时是2.45倍，混合并行+Partial FC时是1.38倍，同时在显存管理和最大支持的类别数上也表现出色。
概述
人脸识别技术可以准确识别出图像中的人脸和身份，具有丰富的应用场景，譬如金融场景下的刷脸支付、安防场景下的罪犯识别和医学场景下的新冠流行病学调查等等。人脸识别的算法演变经历了以 PCA 为代表的早期阶段，再到以“人工特征+分类器”为主的统计学习方法阶段，近几年，随着大数据及 GPU 算力的爆发，人脸识别进入到深度学习算法为绝对主角的阶段。 
人脸识别的规模也从最初基于几百张图片、识别数十张人脸，到现在动辄识别上百万甚至亿级别的人脸，这样超大规模的人脸识别需求，带来了诸多要解决的技术难题，特别是基于分布式机器学习系统训练大规模人脸识别模型的挑战。采用混合并行的方式解决超大规模人脸识别问题，已经成为工业界的最佳实践，但主流深度学习框架（如TensorFlow, PyTorch, MXNet等）仅支持了更容易实现的数据并行，而如果要实现混合并行，往往需要基于深度学习框架进行二次开发且编程复杂度高。
OneFlow 作为主打分布式易用和性能的框架，有着去中心化的Actor机制及SBP的抽象，天然对分布式有着良好支持，训练速度快、显存占用低，同样的显卡能支持更大的batch size，从单机多卡拓展到多机分布式训练并不需要复杂的设置。本文基于OneFlow实现了大规模人脸识别的解决方案，该方案可以帮助用户轻松使用数据并行 + 模型并行的混合训练，同时还支持了Partial FC采样技术，理论上支持过亿(ID数)级别的人脸识别。
InsightFace 是基于 MXNet 框架实现的业界主流人脸识别解决方案。相较MXNet的实现方案，基于OneFlow的实现方案在性能方面更是十分优秀，OneFlow在数据并行时速度是其2.82倍；模型并行时速度是其2.45倍；混合并行+Partial fc时速度是其1.38倍。基于OneFlow实现的代码已合并至 insightface的官方仓库，其中包含了数据集制作教程、训练和验证脚本、预训练模型以及和MXNet模型的转换工具。在 The 1:1 verification accuracy on InsightFace Recognition Test (IFRT) 验证集上，Oneflow及MXNet训练模型的精度对比如下：
| Framework | African | Caucasian | Indian | Asian | All |
|---|---|---|---|---|---|
| OneFlow | 90.4076 | 94.583 | 93.702 | 68.754 | 89.684 |
| MXNet | 90.45 | 94.60 | 93.96 | 63.91 | 88.23 |
仓库地址：https://github.com/deepinsight/insightface/tree/master/recognition/oneflow_face
下面，将以大规模人脸识别任务为例，介绍基于OneFlow和MXNet框架是如何现业界流行的大规模人脸识别技术，以及数据并行、模型并行以及Partial FC采样的实现细节，主要内容包含：
-
1.大规模人脸识别背景介绍 -
2.OneFlow如何实现大规模人脸识别 -
3.MXNet的大规模人脸识别方案 -
4.数据并行、模型并行解决方案的通信量对比 -
5.Partail FC采样技术 -
6.OneFlow和MXNet实现的性能对比 -
7.总结
完全理解本文的技术内容需要对 OneFlow 的并行观和原理有所了解，OneFlow团队曾写过几篇文章详细介绍了OneFlow的并行观、SBP属性等概念，有需要的读者可参考：
1.大规模人脸识别背景介绍
前面的概述简单介绍了人脸识别的一些应用场景以及超大规模人脸识别所带来的一些挑战，下面将更具体一点，来分析下大规模人脸识别方案中涉及到的关键技术以及具体难在何处。
1.1 面临的问题
简单的网络
基于深度学习网络的大规模人脸识别方法一般都是基于常规的 CNN 网络（如resnet50、resnet100等）来提取输入图片中的人脸特征，然后人脸特征会输入到全连接层，最后基于全连接的输出计算 Loss，单纯从网络结构来看，大规模人脸识别相关的网络通常如下图所示： 
复杂的MarginLoss
大规模人脸识别面临的第一个问题是，当使用常规的CNN+FC的网络，后接softmax交叉熵损失函数来对人脸进行分类时，往往得到的结果并不能令人满意，同一个人在不同角度，不同光线等情况下，往往会得到不一样的分类结果。所以在最后计算交叉熵损失之前，需要通过一些方法对FC层的输出进行处理，譬如加入特征间距离的度量，用于判断类别内部（同一个人的不同脸部）和外部（不同人脸间）之间的差异大小。简单来说，一个鲁棒的人脸分类模型应该可以使得映射后的特征具有较小的类内和较大的类间距离。
通过这种基于类间、类内距离度量的方法衍生出了一系列距离度量损失函数，如 center loss、triplet loss、sphereface loss、cosface loss、arcface loss 等等。其中arcface loss源于论文《ArcFace: Additive Angular Margin Loss for Deep Face Recognition》，由于应用了超球面度量、余弦距离等创新方法，其对于类间和类内距离度量更为精准，在各种人脸识别任务和数据集上表现优异。
在 insightface 的源码中，对应 arcface loss 和 cosface loss 的实现是 MarginLoss 类：
class MarginLoss(object):
""" Default is Arcface loss
"""
def __init__(self, margins=(1.0, 0.5, 0.0), loss_s=64, embedding_size=512):
"""
"""
# margins
self.loss_m1 = margins[0]
self.loss_m2 = margins[1]
self.loss_m3 = margins[2]
self.loss_s = loss_s
self.embedding_size = embedding_size
def forward(self, data, weight, mapping_label, depth):
"""
"""
with autograd.record():
norm_data = nd.L2Normalization(data)
norm_weight = nd.L2Normalization(weight)
#
fc7 = nd.dot(norm_data, norm_weight, transpose_b=True)
#
mapping_label_onehot = mx.nd.one_hot(indices=mapping_label,
depth=depth,
on_value=1.0,
off_value=0.0)
# cosface
if self.loss_m1 == 1.0 and self.loss_m2 == 0.0:
_one_hot = mapping_label_onehot * self.loss_m3
fc7 = fc7 - _one_hot
else:
fc7_onehot = fc7 * mapping_label_onehot
cos_t = fc7_onehot
t = nd.arccos(cos_t)
if self.loss_m1 != 1.0:
t = t * self.loss_m1
if self.loss_m2 != 0.0:
t = t + self.loss_m2
margin_cos = nd.cos(t)
if self.loss_m3 != 0.0:
margin_cos = margin_cos - self.loss_m3
margin_fc7 = margin_cos
margin_fc7_onehot = margin_fc7 * mapping_label_onehot
diff = margin_fc7_onehot - fc7_onehot
fc7 = fc7 + diff
fc7 = fc7 * self.loss_s
return fc7, mapping_label_onehot
MarginLoss包含m1~m3这3个参数，通过这3个参数的组合来实现cosface loss和arcface loss。而在oneflow中的实现中则更为简便，直接调用flow.combined_margin_loss即可实现MarginLoss系列的功能。
巨大的人脸ID数导致显存爆炸
对于工业界的人脸识别业务，人脸 ID 数通常会超过百万级，甚至可以达到千万级至亿级别，在这种情况下全连接层的参数矩阵通常会超出单个 GPU 设备的显存上限，所以，仅仅靠普通的数据并行也无法完成训练。而这就是大规模人脸识别方案的核心挑战所在。
1.2 解决方案
数据并行 or 模型并行
为了处理上面提到的问题，工业界对于超大规模的人脸识别任务，往往采用数据并行 + 模型并行的混合并行方式。即在网络前面的CNN部分，采用数据并行进行人脸特征提取，而最后的全连接层则采用模型并行，将参数矩阵切分到多个 GPU 上。
2.OneFlow 如何实现大规模人脸识别
基于OneFlow实现的大规模人脸识别方案对齐了 insightface官方的partail_fc 的实现（基于MNXet），支持数据并行、数据|模型混合并行和Partial FC采样技术，在loss方面支持设置了 m1,m2和m3超参以定义 softmax loss、arcface loss、cosface loss 以及其他组合形式的 combined loss。代码已合并至insightface官方仓库—oneflow_face。
下面将通过整体结构和技术细节实现这两个层面来介绍基于OneFlow的大规模人脸识别方案。
2.1 整体结构
首先是 采用数据并行的CNN 特征提取部分，CNN提取的特征(Features)作为后面的全连接(FC)层的输入，全连接层采用模型并行。全连接层fc1经过Margin loss layer(fc7)处理后的输出，同label一起计算softmax交叉熵损失，得到最终的loss。
整体的网络结构如下图：
FC层以及具体loss计算的细节，如下图所示： 
图中展示了每个GPU设备上具体的计算流程，在GPU上方有全连接层的权重矩阵Weight（图中的matmul节点），黄色长方体表示的Features经CNN提后取的人脸特征。
对于batch_size大小的批量图片输入，Feature的形状为 (batch_size, emb_size)，emb_size根据网络不同通常为128或512。图中的权重(Weight)的大小与人脸类别 ID 数有关，在大规模人脸识别的工业实践中，类别 ID 数通常为百万到亿级别，假设类别 ID 数为1千万，则模型大小为(emb_size, 10000000)。经过全连接层的特征矩阵和权重矩阵相乘后((batch_size, emb_size) × (emb_size, 10000000))的输出特征形状为 (batch_size, 10000000)。
在OneFlow的实现方案中全连接层采用模型并行，即对权重矩阵做切分而使用全量的特征数据，因此输入特征的 SBP 属性为 Broadcast，即表示每个GPU设备上都会拷贝一份特征数据；而权值的 SBP 属性为 Split(1)，即参数矩阵在维度1被切分到各个 GPU 设备上，假设有 P 个 GPU，则每个 GPU 上有 (emb_size, 10000000/P) 大小的权值 。全连接层的输出形状取决于输入features的形状以及类别ID数，故每个设备上的输出形状为 (batch_size, 10000000/P)，且 SBP 属性也为 Split(1)。
由于本方案对权值做了切割(Split(1))，故通常来说，需要对全连接层的输出做合并，并转为按 Split(0) 切分的数据并行，但是由于全连接层的输出数据块较大，如果直接由 Split(1) 转为 常规的Split(0)，会引入大量(可以避免的)通信。因此，在 OneFlow 算子实现的内部，并不先进行 Split 的转化，而是将全连接层的输出直接作为 softmax 的输入进行计算，因为 softmax 的运算特性，可以使得输出的 SBP 属性依然是 Split(1)。类似的，softmax 的输出（按 Split(1) 切分）继续作为 sparse_cross_entropy 的输入进行计算，由于算子本身的特性和 OneFlow 的机制，sparse_cross_entropy 的输出的 SBP 属性依然可以保持 Split(1)。
经过 sparse_cross_entropy 处理后的输出，逻辑上获得最终的 loss 结果。 此时，数据块的形状为 (batch_size, 1)，已经很小，这时候再将模型并行的Split(1）模式转为按 Split(0) 切分的数据并行。
以上即是 OneFlow 实现的全连接层的内部工作流程，对于普通算法开发者来说，了解以上内容即掌握了OneFlow大规模人脸方案中全连接层处理的核心流程。对于框架开发者和实现细节感兴趣的朋友，请看下面的小节—2.全连接层的技术细节。这一小结，将会用较大篇幅展开softmax和sparse_cross_entropy实现细节相关的内容以及Split(1）切分是如果做到数学上的等价。
2.2 Oneflow实现代码解析
下面，我们讲解一下在Oneflow是如何通过简单的几行代码来实现大规模人脸识别方案。首先，backbone部分的网络是类似的由CNN+FC全连接层构成，我们重点看一下FC之后的Marginloss层及相关处理。主要代码如下：
elif config.loss_name == "margin_softmax":
if args.model_parallel:
print("Training is using model parallelism now.")
labels = labels.with_distribute(flow.distribute.broadcast())
fc1_distribute = flow.distribute.broadcast()
fc7_data_distribute = flow.distribute.split(1)
fc7_model_distribute = flow.distribute.split(0)
else:
fc1_distribute = flow.distribute.split(0)
fc7_data_distribute = flow.distribute.split(0)
fc7_model_distribute = flow.distribute.broadcast()
fc7_weight = flow.get_variable(
name="fc7-weight",
shape=(config.num_classes, embedding.shape[1]),
dtype=embedding.dtype,
initializer=_get_initializer(),
regularizer=None,
trainable=trainable,
model_name="weight",
distribute=fc7_model_distribute,
)
if args.partial_fc and args.model_parallel:
print(
"Training is using model parallelism and optimized by partial_fc now."
)
(
mapped_label,
sampled_label,
sampled_weight,
) = flow.distributed_partial_fc_sample(
weight=fc7_weight, label=labels, num_sample=args.total_num_sample,
)
labels = mapped_label
fc7_weight = sampled_weight
fc7_weight = flow.math.l2_normalize(
input=fc7_weight, axis=1, epsilon=1e-10)
fc1 = flow.math.l2_normalize(
input=embedding, axis=1, epsilon=1e-10)
fc7 = flow.matmul(
a=fc1.with_distribute(fc1_distribute), b=fc7_weight, transpose_b=True
)
fc7 = fc7.with_distribute(fc7_data_distribute)
fc7 = (
flow.combined_margin_loss(
fc7, labels, m1=config.loss_m1, m2=config.loss_m2, m3=config.loss_m3
最低0.47元/天 解锁文章
扫一扫
895
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
