用pytorch实现arcface loss,从训练到部署(1)
目标和思路
本人是深度学习的小白一枚,以前一直用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的遇到的坑记录一下,以便绕开。