Set MobileNetV3 `feature_maps` parameter (#515)

configs/ssd/ssdlite_mobilenet_v3_large.yml

@@ -26,8 +26,8 @@ MobileNetV3:
   scale: 1.0
   model_name: large
   extra_block_filters: [[256, 512], [128, 256], [128, 256], [64, 128]]
-  with_extra_blocks: true
   conv_decay: 0.00004
+  feature_maps: [5, 7, 8, 9, 10, 11]
 
 SSDLiteMultiBoxHead:
   aspect_ratios: [[2.], [2., 3.], [2., 3.], [2., 3.], [2., 3.], [2., 3.]]

configs/ssd/ssdlite_mobilenet_v3_small.yml

@@ -26,8 +26,8 @@ MobileNetV3:
   scale: 1.0
   model_name: small
   extra_block_filters: [[256, 512], [128, 256], [128, 256], [64, 128]]
-  with_extra_blocks: true
   conv_decay: 0.00004
+  feature_maps: [5, 7, 8, 9, 10, 11]
 
 SSDLiteMultiBoxHead:
   aspect_ratios: [[2.], [2., 3.], [2., 3.], [2., 3.], [2., 3.], [2., 3.]]

configs/yolov3_mobilenet_v3.yml

@@ -19,7 +19,7 @@ MobileNetV3:
   norm_decay: 0.
   model_name: large
   scale: 1.
-  with_extra_blocks: false
+  feature_maps: [1, 2, 3, 4, 6]
 
 YOLOv3Head:
   anchor_masks: [[6, 7, 8], [3, 4, 5], [0, 1, 2]]

ppdet/modeling/backbones/mobilenet_v3.py

@@ -12,12 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import OrderedDict
+
 import paddle.fluid as fluid
 from paddle.fluid.param_attr import ParamAttr
 from paddle.fluid.regularizer import L2Decay
 
 from ppdet.core.workspace import register
-import math
+from numbers import Integral
 
 __all__ = ['MobileNetV3']
 
@@ -32,7 +38,7 @@ class MobileNetV3():
         norm_type (str): normalization type, 'bn' and 'sync_bn' are supported.
         norm_decay (float): weight decay for normalization layer weights.
         conv_decay (float): weight decay for convolution layer weights.
-        with_extra_blocks (bool): if extra blocks should be added.
+        feature_maps (list): index of stages whose feature maps are returned.
         extra_block_filters (list): number of filter for each extra block.
     """
     __shared__ = ['norm_type']
@@ -40,21 +46,24 @@ class MobileNetV3():
     def __init__(self,
                  scale=1.0,
                  model_name='small',
-                 with_extra_blocks=False,
+                 feature_maps=[5, 6, 7, 8, 9, 10],
                  conv_decay=0.0,
                  norm_type='bn',
                  norm_decay=0.0,
                  extra_block_filters=[[256, 512], [128, 256], [128, 256],
                                       [64, 128]]):
+        if isinstance(feature_maps, Integral):
+            feature_maps = [feature_maps]
+
         self.scale = scale
         self.model_name = model_name
-        self.with_extra_blocks = with_extra_blocks
+        self.feature_maps = feature_maps
         self.extra_block_filters = extra_block_filters
         self.conv_decay = conv_decay
         self.norm_decay = norm_decay
         self.inplanes = 16
         self.end_points = []
-        self.block_stride = 1
+        self.block_stride = 0
         if model_name == "large":
             self.cfg = [
                 # kernel_size, expand, channel, se_block, act_mode, stride
@@ -181,8 +190,11 @@ class MobileNetV3():
             if_act=True,
             act=act,
             name=name + '_expand')
-        if self.block_stride == 16 and stride == 2:
-            self.end_points.append(conv0)
+        if self.block_stride == 4 and stride == 2:
+            self.block_stride += 1
+            if self.block_stride in self.feature_maps:
+                self.end_points.append(conv0)
+
         conv1 = self._conv_bn_layer(
             input=conv0,
             filter_size=filter_size,
@@ -265,9 +277,11 @@ class MobileNetV3():
             name='conv1')
         i = 0
         for layer_cfg in cfg:
-            self.block_stride *= layer_cfg[5]
             if layer_cfg[5] == 2:
-                blocks.append(conv)
+                self.block_stride += 1
+                if self.block_stride in self.feature_maps:
+                    self.end_points.append(conv)
+
             conv = self._residual_unit(
                 input=conv,
                 num_in_filter=inplanes,
@@ -280,10 +294,9 @@ class MobileNetV3():
                 name='conv' + str(i + 2))
             inplanes = int(scale * layer_cfg[2])
             i += 1
-        blocks.append(conv)
-
-        if not self.with_extra_blocks:
-            return blocks
+        self.block_stride += 1
+        if self.block_stride in self.feature_maps:
+            self.end_points.append(conv)
 
         # extra block
         conv_extra = self._conv_bn_layer(
@@ -296,13 +309,18 @@ class MobileNetV3():
             if_act=True,
             act='hard_swish',
             name='conv' + str(i + 2))
-        self.end_points.append(conv_extra)
+        self.block_stride += 1
+        if self.block_stride in self.feature_maps:
+            self.end_points.append(conv_extra)
         i += 1
         for block_filter in self.extra_block_filters:
             conv_extra = self._extra_block_dw(conv_extra, block_filter[0],
                                               block_filter[1], 2,
                                               'conv' + str(i + 2))
-            self.end_points.append(conv_extra)
+            self.block_stride += 1
+            if self.block_stride in self.feature_maps:
+                self.end_points.append(conv_extra)
             i += 1
 
-        return self.end_points
+        return OrderedDict([('mbv3_{}'.format(idx), feat)
+                            for idx, feat in enumerate(self.end_points)])