用pytorch实现arcface loss，从训练到部署

目标和思路

本人是深度学习的小白一枚，以前一直用TensorFlow玩自己的数据集，最近突发奇想，想真正搞一个图像检索应用（100万级别的数据量，15万类），把训练和部署等流程走一遍。

为实现这个目标，思路如下：
**第一步。**肯定是选择一个合适CNN作为Backbone提取图像特征，这里选择很多，比如：Inception、Resnet、MobileNet等等。以MobileNet为例，在TensorFlow2.0下，直接一行代码就可以加载内设的模型：

tensorflow.keras.applications.MobileNetV2()

看起来十分方便。

**第二步。**就是在Backbone后面追加人脸识别中经常用到的cosine loss、arcface loss等自定义层。
为什么追加？因为Backbone一般输出1000类（ImageNet基准是1000类），在我们的实践中，15万类远大于1000类，就需要对softmax函数进行改造，把Backbone生成的特征feature映射到角度空间内，让不同类之间的距离尽可能大，想同类之间的距离尽可能小（白话翻译，理论公式自行百度）。cosine loss、arcface loss等都是这个目标。
怎么追加？最简单的做法，不管Backbone最终输出的特征有多大，直接先追加一层全连接层，把特征统一到512维，当然也可以是256维或者更大，具体根据数据集的不同进行调整。然后将512维的特征值和label标签一期作为cosine loss、arcface loss层的输入，然后计算交叉熵Loss。
当然，增加全连接层会导致模型参数增加。可以根据每个Backbone的输出不同，使用池化等操作减少参数。

**第三步。**把Backbone和arcface loss层组装成一个模型进行训练。

**第四步。**搞个服务器部署模型。

为什么选择PyTorch

TensorFlow在工业界的应用十分广泛。因为他有非常完备的工业应用支持。尤其是TensorFlow Serving，可以很方便地实现模型部署，不停机升级等等功能。但TensorFlow学习难度大，尤其是1.x版本的静态图机制简直就是个坑。而PyTorch在学术界应用较多，不足之处就是工业应用支持不足。

但是，随着微软和Facebook联合发布的onnx热度飙升，PyTorch可以很方便地部署模型，也可以通过onnx这座桥梁，让模型被TensorFlow、caffe2、mxnet、paddlepaddle等框架调用。此外，onnx还开源了一个onnx runtime server，可以直接提供REST API服务。

-------重点来了—>

在训练包含arcface loss层的模型时，TensorFlow一直出现loss 为NAN的问题。各种搜索引擎、各种论坛都求助了，但是一直没有找到好的办法。
有的说是tf.sqrt这个函数在开平方根的时候遇到0值会导致NAN，增加了1e-8方式0值，也没用。
有的说是学习率、margin超参和scale超参的初始值不对，各种尝试之后也不行。
把Pytorch下实现的arcface层照搬到TensorFlow下，还是不行。

心灰意冷，投靠PyTorch。

用了PyTorch后，发现他还有个实用的地方，就是可以将Backbone和arcface的模型分开保存，在部署的时，只用到Backbone的模型，运行得到feature后，直接用紧邻搜索库计算相似度，岂不是很完美。

下一步，就把我用pytorch的遇到的坑记录一下，以便绕开。