tensorflow object detect API网络结构源码修改

最近使用TensorFlow object_detect API做目标检测任务,由于要求目标检测模型能够移植客户端中,进而选择目标检测模型时则选择轻量级的模型,最后选择了ssd_mobilenet_v1作为目标检测的模型。

之前写过了TensorFlow object_detect API训练自己数据的步骤以及通过修改配置文件参数降低模型输入大小和模型通道数方法达到压缩模型提高推理速率(此处)。本篇文章主要讲述如何通过修改object_detect API中网络结构源码,以ssd_mobilenet_v1为例。

  • 1.模型结构定义源码

若要在源码中修改mobilenet-ssd-v1结构,可修改的源码文件:

    (1). models-master/research/slim/nets/mobilenet_v1.py   #模型网络的主结构

    (2). models-master/research/object_detection/models/ssd_mobilenet_v1_feature_extractor.py     #anchors特征提取层

mobilenet_v1.py:

tensorflow object detect API网络结构源码修改

 

mobilenet_v1.py主要定义ssd_mobilenet_v1网络的主结构层,若需要对网络进行删减可修改上图中MOBILENET_CONV_DEFS变量。(若进行网络层删减后配置文件中的层数也需要进行相应修改)

ssd_mobilenet_v1_feature_extractor.py:

tensorflow object detect API网络结构源码修改

 

ssd_mobilenet_v1_feature_extractor.py为anchors特征提取层,anchors特征提取对应的层修改可以修改上图中的feature_map_layout变量。(anchors特征层删减后配置文件中的相关参数也需要进行修改)

配置文件参数如下:

tensorflow object detect API网络结构源码修改

 

  • 2. anchors特征提取层修改实践

接下来举例如何通过修改源码文件ssd_mobilenet_v1_feature_extractor.py将anchors提取的6层结构改为4层。

首先修改源码文件

models-master/research/object_detection/models/ssd_mobilenet_v1_feature_extractor.py,将feature_map_layout变量中的from_layer和layer_depth后2层去掉如下:

tensorflow object detect API网络结构源码修改

 

然后修改配置文件如下:

将num_layers由6改为4,并去掉2个aspect_ratios

tensorflow object detect API网络结构源码修改

 

这样便可以开始目标检测模型训练(由于网络结构修改了,预加载官方模型的参数需要在配置文件中去掉,改用不预加载模型的方法训练)

 

  • 3. 网络结构修改前后对比图

修改前网络结构(部分结构):

tensorflow object detect API网络结构源码修改

 

修改后网络结构(部分结构):

tensorflow object detect API网络结构源码修改

 

    可从修改前后的两个网络结构图中看出通过源码修改后,目标检测最后6层的特征提取成功变为4层,总体上能够减少一些计算量(后面的层计算量很少)。4个anchors特征提取层的目标检测模型依旧能够正常检测(精度会有相应损失)。