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

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 194,242评论 5 459
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 81,769评论 2 371
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 141,484评论 0 319
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 52,133评论 1 263
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 61,007评论 4 355
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 46,080评论 1 272
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 36,496评论 3 381
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 35,190评论 0 253
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 39,464评论 1 290
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 34,549评论 2 309
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 36,330评论 1 326
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 32,205评论 3 312
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 37,567评论 3 298
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 28,889评论 0 17
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 30,160评论 1 250
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 41,475评论 2 341
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 40,650评论 2 335