MobileNet SSD V2模型的压缩与tflite格式的转换

最近项目里需要一个小型的目标检测模型,SSD、YOLO等一通模型调参试下来,直接调用TensorFlow object detect API居然效果最好,大厂的产品不得不服啊。使用mobilenet ssd v2模型,配置文件也未修改参数,训练后的模型不光检测效果不错,在CPU上的运行时间也在70ms左右。之后将模型移植到安卓手机上(魅族MX4,老的不是一点点),卡顿明显;改用同事的华为,在麒麟960上略微流畅了一些,但仍然不能达到实时检测。而且训练得到的pb模型居然有19M,实在太大了,于是又探索了一波模型的压缩和量化。

模型压缩

说到模型压缩,最简单粗暴的方法当然是减少卷积层数。在使用Tensorflow的API之前,我训练过一个SSD模型,检测效果不错,但耗时接近1s。为了提高检测速度我果断开始减少卷积层数,并做了不同层数的对比试验。结果和原始的VGG16骨干相比,要么检测效果相近,耗时也没少多少,要么耗时大减,但漏检率飙升。也就是在这个情况下,我转投了mobilenet网络。

所以这次面临模型压缩时, 我没有再尝试这个选项(当然也有配置文件不支持删减层数,要删就要去改slim里的源码这个原因。我一个前同事是中科院计算机博士,他的格言就是觉得源码不好就别调用,自己写;要调用就尽量避免改源码,因为你肯定没有源码写得好)。这样看下来,就只能在配置文件的范围内自由发挥了。

修改配置文件

首先,附上Tensorflow object detection API中支持的各大模型的配置文件地址:

models/research/object_detection/samples/configs at master · tensorflow/models · GitHub

这里面关于mobilenet_ssd_v2的有好几个:


我使用的是最经典的基于COCO数据集训练的配置文件,也就是第一个。图里的最后一个也是基于COCO数据集的,不过是有量化的模型,这个文件我在后面也有用到。

打开配置文件,里面主要分成model、train和eval三块。在调用API训练自己的数据时,train和eval的数据当然是要修改的:


回到model部分,在feature_extractor那里,有一个depth_multiplier,这个参数作为一个因子与网络中各层的channel数相乘,换言之,depth_multiplier越小,网络中feature map的channel数越少,模型参数自然也就少了很多。depth_multiplier默认为1,在我的实验里改成了0.25,试就试一把大的。

训练模型

之前depth_multiplier为1时, 我训练是加载了预训练模型的,模型地址:

models/detection_model_zoo.md at master · tensorflow/models · GitHub


从图中可以看出,mobilenet_v1的预训练模型中有一种0.75_depth的版本,这就是depth_multiplier取0.75时在COCO数据集上训练出来的模型。对于mobilenet_v2,只提供了非量化版和量化版(个人觉得应该0.25、0.5、0.75这几个常用档都提供一个,难道是官方不建议压缩太多吗。。。)

由于没有对应的预训练模型,所以可以选择加载或者不加载模型。

加载模型的话,开始训练后命令行会打印一大堆XXX is available in checkpoint, but has an incompatible shape with model variable. This variable will not be initialized from the checkpoint. 不过这并不影响训练,忽略就可以了。

不加载的话,就将配置文件里fine_tune_checkpoint的那两行注释掉。

进入到object detection目录,运行python object_detection/model_main.py  --pipeline_config_path=xxxxxxx/ssd_mobilenet_v2_coco.config  --model_dir=xxxxxxxx即可

PS:训练过程中是不会打印训练信息的,看命令行会以为电脑卡住了。。。直到eval才会打印出信息

PPS:可以通过TensorBoard来监听训练过程,判断训练是在正常进行还是电脑真的卡住了(这种情况可能是因为batch size和输入图片大小太大。默认是24和300*300,但也都可以改)

模型导出

训练完成之后,还是在object detection目录下,运行python export_inference_graph.py,必要的参数分别是输入的ckpt的文件地址,输出的pb文件的文件夹以及配置文件地址。

在深度压缩至0.25倍之后, 我的pb模型大小仅为2.2M,效果卓群。当然网络的缩减会带来精度的损失,我的AR和AP分别降了2个点和3个点。

模型移植

生成tflite模型

Tensorflow object detection API训练出的模型,讲道理从ckpt转成tflite只需要两步:

第一步,将ckpt转成pb文件,这次使用的是python export_tflite_ssd_graph.py,操作难度不大,会得到tflite_graph.pb和tflite_graph.pbtxt两个文件;

第二步,将pb转为tflite文件,我搜到的方法大都是使用bazel编译tensorflow/contirb/lite/toco下面的toca文件,但我反复尝试,报了多种错误,依旧没有成功。。。最后我在stackoverflow上搜到了一位小哥的回复,进入tensorflow/contrib/lite/python目录,运行python tflite_convert.py,参数设置为

--graph_def_file=XXX/tflite_graph.pb 上一步生成的pb文件地址

--output_file=XXX/xxx.tflite 输出的tflite文件地址

--input_arrays=normalized_input_image_tensor 输入输出的数组名称对于mobilenet ssd是固定的,不用改 

--output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'

 --input_shape=1,XXX,XXX,3 输入的图片大小,需要与配置文件中一致

--allow_custom_ops

验证tflite模型

在将tflite模型放进手机之前,我在python里加载tflite模型测试了一次,流程类似加载pb模型

第一步,导入模型

interpreter = tf.contrib.lite.Interpreter(model_path="compress_export/detect.tflite")

interpreter.allocate_tensors()

第二步,获得输入和输出的tensor

input_details = interpreter.get_input_details()

output_details = interpreter.get_output_details()

第三步,读取输入图像,feed给输入tensor

可以采用PIL或cv2将图像读入,转为numpy数组,然后赋值给input_data

input_data = np.array(XXX)

interpreter.set_tensor(input_details[0]['index'], input_data)

第四步,运行模型

interpreter.invoke()

第五步, 获得输出

参考输入tensor的表示方法,目标检测的输出有4个,具体的值可以通过output_details[0]['index']、output_details[1]['index']、output_details[2]['index']、output_details[3]['index']获得

这里有一个我踩到的坑,验证tflite模型时,我采用了和加载pb模型完全相同的图片预处理步骤,输出的结果完全不同。几番检查之后,发现问题出在模型转换时。运行python tflite_convert.py时,输入数组的名称为normalized_input_image_tensor,而我训练时采用的是未经normalized的数组。所以在模型转换时,tensorflow内置了对input进行normalized的步骤。因此在调用tflite模型时,同样需要在图像预处理中加入这一步。 nomlized的方法为除以128.0再减去1,保证输入的值在[-1,1)范围内。

参考

https://blog.csdn.net/qq_26535271/article/details/84930868

Tensorflow Convert pb file to TFLITE using python - Stack Overflow