MobileNet SSD V2模型的压缩与tflite格式的转换（补充版）

Posted 2023-03-23

参考技术A 最近项目里需要一个小型的目标检测模型，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个点。

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模型放进手机之前，我在python里加载tflite模型测试了一次，流程类似加载pb模型

第一步，导入模型

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

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

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

第四步，运行模型

第五步， 获得输出

参考输入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)范围内。

尝试向检测模型“ssd_mobilenet_v2”添加警报，会引发错误

【中文标题】尝试向检测模型“ssd_mobilenet_v2”添加警报，会引发错误

【英文标题】：Trying to add an alert to detection model "ssd_mobilenet_v2", throws an error

【发布时间】：2022-01-09 11:40:00

【问题描述】：

我正在尝试使用“ssd_mobilenet_v2_fpn_keras”添加警报系统

检测模型加载到下面的函数中

 def detect_fn(image):
    image, shapes = detection_model.preprocess(image)
    prediction_dict = detection_model.predict(image, shapes)
    detections = detection_model.postprocess(prediction_dict, shapes)
    return detections

图像转换为张量

input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)

张量被馈送到检测模型

detections = detect_fn(input_tensor)

检测模型的输出是一个字典，具有以下键：

dict_keys(['detection_boxes', 'detection_scores', 'detection_classes', 'raw_detection_boxes', 'raw_detection_scores', 'detection_multiclass_scores', 'detection_anchor_indices', 'num_detections'])

detections[detection_classes]，给出以下输出，即 0 是 ClassA，1 是 ClassB

[0 1 1 0 0 1 0 0 1 0 1 1 0 0 1 0 1 1 0 1 0 1 1 0 0 1 0 0 1 0 1 0 0 1 1 1 1 0 0 0 1 1 1 0 0 1 1 1 0 1 0 1 0 0 0 0 1 0 0 1 0 0 1 0 1 0 0 1 0 0 0 0 1 0 1 1 0 1 1 0 1 1 0 0 0 0 0 0 0 0 0 1 1 1 0 0 0 1 0 1]

detections['detection_scores'] 给出检测到的每个框的分数（下面显示了几个）

[0.988446 0.7998712 0.1579772 0.13801616 0.13227147 0.12731305 0.09515342 0.09203091 0.09191579 0.08860824 0.08313078 0.07684237

我正在尝试Print("Attention needed")，如果观察到检测类 B 即 1

for key in detections['detection_classes']:
if key==1:
    print('Alert')

当我尝试这样做时，我得到一个错误

`ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

如何让它发挥作用？

我希望代码打印“需要注意”是 Class =1 或 A 并且 detection_scores >= 14

Code Explained, a bit further

---

完整代码的链接如下：

Tutorial on YouTube GitHub sources repository

【参考方案1】：

如错误消息中所述，您应该使用.any()。喜欢：

if (key == 1).any():
  print('Alert')

因为key == 1 将是一个带有[False, True, True, False, ...] 的数组

您可能还想检测超过特定分数的分数，例如 0.7：

for key, score in zip(
  detections['detection_classes'],
  detections['detection_scores']):
  if score > 0.7 and key == 1:
    print('Alert')
    break