update

paddleslim/nas/search_space/inception_block.py

@@ -22,7 +22,7 @@ from paddle.fluid.param_attr import ParamAttr
 from .search_space_base import SearchSpaceBase
 from .base_layer import conv_bn_layer
 from .search_space_registry import SEARCHSPACE
-from .utils import compute_downsample_num
+from .utils import compute_downsample_num, check_points
 
 __all__ = ["InceptionABlockSpace", "InceptionCBlockSpace"]
 ### TODO add asymmetric kernel of conv when paddle-lite support 
@@ -162,9 +162,7 @@ class InceptionABlockSpace(SearchSpaceBase):
                          self.k_size[tokens[i * 9 + 7]], 1,
                          self.pool_type[tokens[i * 9 + 8]]))
 
-        def net_arch(input, return_mid_layer=False, return_block=[]):
-            assert isinstance(return_block,
-                              list), 'return_block must be a list.'
+        def net_arch(input, return_mid_layer=False, return_block=None):
             layer_count = 0
             mid_layer = dict()
             for i, layer_setting in enumerate(self.bottleneck_params_list):
@@ -174,7 +172,7 @@ class InceptionABlockSpace(SearchSpaceBase):
                 pool_type = 'avg' if layer_setting[9] == 0 else 'max'
                 if stride == 2:
                     layer_count += 1
-                if (layer_count - 1) in return_block:
+                if check_points((layer_count - 1), return_block):
                     mid_layer[layer_count - 1] = input
 
                 input = self._inceptionA(
@@ -406,9 +404,7 @@ class InceptionCBlockSpace(SearchSpaceBase):
                          self.k_size[tokens[i * 11 + 9]], 1,
                          self.pool_type[tokens[i * 11 + 10]]))
 
-        def net_arch(input, return_mid_layer=False, return_block=[]):
-            assert isinstance(return_block,
-                              list), 'return_block must be a list.'
+        def net_arch(input, return_mid_layer=False, return_block=None):
             layer_count = 0
             mid_layer = dict()
             for i, layer_setting in enumerate(self.bottleneck_params_list):
@@ -418,7 +414,7 @@ class InceptionCBlockSpace(SearchSpaceBase):
                 pool_type = 'avg' if layer_setting[11] == 0 else 'max'
                 if stride == 2:
                     layer_count += 1
-                if (layer_count - 1) in return_block:
+                if check_points((layer_count - 1) in return_block):
                     mid_layer[layer_count - 1] = input
 
                 input = self._inceptionC(

paddleslim/nas/search_space/mobilenet_block.py

@@ -22,7 +22,7 @@ from paddle.fluid.param_attr import ParamAttr
 from .search_space_base import SearchSpaceBase
 from .base_layer import conv_bn_layer
 from .search_space_registry import SEARCHSPACE
-from .utils import compute_downsample_num
+from .utils import compute_downsample_num, check_points
 
 __all__ = ["MobileNetV1BlockSpace", "MobileNetV2BlockSpace"]
 
@@ -133,9 +133,7 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
                          self.repeat[tokens[i * 4 + 2]], 1,
                          self.k_size[tokens[i * 4 + 3]]))
 
-        def net_arch(input, return_mid_layer=False, return_block=[]):
-            assert isinstance(return_block,
-                              list), 'return_block must be a list.'
+        def net_arch(input, return_mid_layer=False, return_block=None):
             # all padding is 'SAME' in the conv2d, can compute the actual padding automatic. 
             # bottleneck sequences
             in_c = int(32 * self.scale)
@@ -148,7 +146,7 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
 
                 if s == 2:
                     layer_count += 1
-                if (layer_count - 1) in return_block:
+                if check_points((layer_count - 1), return_block):
                     mid_layer[layer_count - 1] = depthwise_conv
 
                 input, depthwise_conv = self._invresi_blocks(
@@ -162,7 +160,7 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
                     name='mobilenetv2_' + str(i + 1))
                 in_c = int(c * self.scale)
 
-            if layer_count in return_block:
+            if check_points(layer_count, return_block):
                 mid_layer[layer_count] = depthwise_conv
 
             if return_mid_layer:
@@ -374,9 +372,7 @@ class MobileNetV1BlockSpace(SearchSpaceBase):
                          self.filter_num[tokens[i * 3 + 1]], 1,
                          self.k_size[tokens[i * 3 + 2]]))
 
-        def net_arch(input, return_mid_layer=False, return_block=[]):
-            assert isinstance(return_block,
-                              list), 'return_block must be a list.'
+        def net_arch(input, return_mid_layer=False, return_block=None):
             mid_layer = dict()
             layer_count = 0
 
@@ -384,7 +380,7 @@ class MobileNetV1BlockSpace(SearchSpaceBase):
                 filter_num1, filter_num2, stride, kernel_size = layer_setting
                 if stride == 2:
                     layer_count += 1
-                if (layer_count - 1) in return_block:
+                if check_points((layer_count - 1), return_block):
                     mid_layer[layer_count - 1] = input
 
                 input = self._depthwise_separable(

paddleslim/nas/search_space/mobilenetv1.py

@@ -22,6 +22,7 @@ from paddle.fluid.param_attr import ParamAttr
 from .search_space_base import SearchSpaceBase
 from .base_layer import conv_bn_layer
 from .search_space_registry import SEARCHSPACE
+from .utils import check_points
 
 __all__ = ["MobileNetV1Space"]
 
@@ -157,7 +158,7 @@ class MobileNetV1Space(SearchSpaceBase):
 
         def net_arch(input,
                      scale=1.0,
-                     return_block=[],
+                     return_block=None,
                      end_points=None,
                      output_stride=None):
             self.scale = scale
@@ -165,15 +166,6 @@ class MobileNetV1Space(SearchSpaceBase):
 
             decode_ends = dict()
 
-            def check_points(count, points):
-                if points is None:
-                    return False
-                else:
-                    if isinstance(points, list):
-                        return (True if count in points else False)
-                    else:
-                        return (True if count == points else False)
-
             input = conv_bn_layer(
                 input=input,
                 filter_size=3,

paddleslim/nas/search_space/mobilenetv2.py

@@ -22,6 +22,7 @@ from paddle.fluid.param_attr import ParamAttr
 from .search_space_base import SearchSpaceBase
 from .base_layer import conv_bn_layer
 from .search_space_registry import SEARCHSPACE
+from .utils import check_points
 
 __all__ = ["MobileNetV2Space"]
 
@@ -136,7 +137,7 @@ class MobileNetV2Space(SearchSpaceBase):
 
         def net_arch(input,
                      scale=1.0,
-                     return_block=[],
+                     return_block=None,
                      end_points=None,
                      output_stride=None):
             self.scale = scale

paddleslim/nas/search_space/resnet.py

@@ -22,17 +22,14 @@ from paddle.fluid.param_attr import ParamAttr
 from .search_space_base import SearchSpaceBase
 from .base_layer import conv_bn_layer
 from .search_space_registry import SEARCHSPACE
+from .util import check_points
 
 __all__ = ["ResNetSpace"]
 
 
 @SEARCHSPACE.register
 class ResNetSpace(SearchSpaceBase):
-    def __init__(self,
-                 input_size,
-                 output_size,
-                 block_num,
-                 block_mask=None):
+    def __init__(self, input_size, output_size, block_num, block_mask=None):
         super(ResNetSpace, self).__init__(input_size, output_size, block_num,
                                           block_mask)
         # self.filter_num1 ~ self.filter_num4 means convolution channel
@@ -91,7 +88,9 @@ class ResNetSpace(SearchSpaceBase):
         num_filters.append(filter4)
         depth.append(repeat4)
 
-        def net_arch(input):
+        def net_arch(input, return_block=None, end_points=None):
+            decode_ends = dict()
+
             conv = conv_bn_layer(
                 input=input,
                 filter_size=5,
@@ -99,14 +98,26 @@ class ResNetSpace(SearchSpaceBase):
                 stride=2,
                 act='relu',
                 name='resnet_conv0')
+            layer_count = 1
             for block in range(len(depth)):
                 for i in range(depth[block]):
+                    stride = 2 if i == 0 and block != 0 else 1
+                    if stride == 2:
+                        layer_count += 1
+                    if check_points((layer_count - 1), return_block):
+                        decode_ends[layer_count - 1] = conv
+
+                    if check_points((layer_count - 1), end_points):
+                        return conv, decode_ends
+
                     conv = self._bottleneck_block(
                         input=conv,
                         num_filters=num_filters[block],
-                        stride=2 if i == 0 and block != 0 else 1,
+                        stride=stride,
                         name='resnet_depth{}_block{}'.format(i, block))
 
+            if check_points(layer_count, end_points):
+                return conv, decode_ends
 
             return conv
 

paddleslim/nas/search_space/resnet_block.py

@@ -22,7 +22,7 @@ from paddle.fluid.param_attr import ParamAttr
 from .search_space_base import SearchSpaceBase
 from .base_layer import conv_bn_layer
 from .search_space_registry import SEARCHSPACE
-from .utils import compute_downsample_num
+from .utils import compute_downsample_num, check_points
 
 __all__ = ["ResNetBlockSpace"]
 
@@ -121,16 +121,14 @@ class ResNetBlockSpace(SearchSpaceBase):
                          self.repeat[tokens[i * 6 + 4]],
                          self.repeat[tokens[i * 6 + 5]], 1))
 
-        def net_arch(input, return_mid_layer=False, return_block=[]):
-            assert isinstance(return_block,
-                              list), 'return_block must be a list.'
+        def net_arch(input, return_mid_layer=False, return_block=None):
             layer_count = 0
             mid_layer = dict()
             for i, layer_setting in enumerate(self.bottleneck_params_list):
                 filter_num1, filter_num2, filter_num3, k_size, repeat1, repeat2, stride = layer_setting
                 if stride == 2:
                     layer_count += 1
-                if (layer_count - 1) in return_block:
+                if check_points((layer_count - 1), return_block):
                     mid_layer[layer_count - 1] = input
 
                 input = self._bottleneck_block(

paddleslim/nas/search_space/utils.py

@@ -14,6 +14,7 @@
 
 import math
 
+
 def compute_downsample_num(input_size, output_size):
     downsample_num = 0
     while input_size > output_size:
@@ -21,6 +22,17 @@ def compute_downsample_num(input_size, output_size):
         downsample_num += 1
 
     if input_size != output_size:
-        raise NotImplementedError('output_size must can downsample by input_size!!!')
+        raise NotImplementedError(
+            'output_size must can downsample by input_size!!!')
 
     return downsample_num
+
+
+def check_points(count, points):
+    if points is None:
+        return False
+    else:
+        if isinstance(points, list):
+            return (True if count in points else False)
+        else:
+            return (True if count == points else False)