update

paddleslim/nas/search_space/inception_block.py

+# Copyright (c) 2019  PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# 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
+
+import numpy as np
+import paddle.fluid as fluid
+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
+
+__all__ = [
+    "InceptionABlockSpace", "InceptionBBlockSpace", "InceptionCBlockSpace"
+]
+### TODO add asymmetric kernel of conv when paddle-lite support 
+
+
+@SEARCHSPACE.register
+class InceptionABlockSpace(SearchSpaceBase):
+    def __init__(self, input_size, output_size, block_num, block_mask):
+        super(InceptionABlockSpace, self).__init__(input_size, output_size,
+                                                   block_num, block_mask)
+        if self.block_mask == None:
+            # use input_size and output_size to compute self.downsample_num
+            self.downsample_num = compute_downsample_num(self.input_size,
+                                                         self.output_size)
+        if self.block_num != None:
+            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(
+                self.downsample_num, self.block_num)
+
+        ### self.filter_num means filter nums
+        self.filter_num = np.array([
+            3, 4, 8, 12, 16, 24, 32, 48, 64, 80, 96, 128, 144, 160, 192, 224,
+            256, 320, 384, 480, 512, 1024
+        ])
+        ### self.k_size means kernel_size
+        self.k_size = np.array([3, 5])
+        ### self.pool_type means pool type, 0 means avg, 1 means max
+        self.pool_type = np.array([0, 1])
+        ### self.repeat means repeat of 1x1 conv in branch of inception
+        ### self.repeat = np.array([0,1])
+
+    def init_tokens(self):
+        """
+        The initial token.
+        """
+        if self.block_mask != None:
+            return [0] * (len(self.block_mask) * 9)
+        else:
+            return [0] * (self.block_num * 9)
+
+    def range_table(self):
+        """
+        Get range table of current search space, constrains the range of tokens.
+        """
+        range_table_base = []
+        if self.block_mask != None:
+            range_table_length = len(self.block_mask)
+        else:
+            range_table_length = self.block_mum
+
+        for i in range(range_table_length):
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.k_size))
+            range_table_base.append(len(self.pooltype))
+
+        return range_table_base
+
+    def token2arch(self, tokens=None):
+        """
+        return net_arch function
+        """
+        #assert self.block_num
+        if tokens is None:
+            tokens = self.init_tokens()
+
+        self.bottleneck_params_list = []
+        if self.block_mask != None:
+            for i in range(len(self.block_mask)):
+                self.bottleneck_params_list.append(
+                    (self.filter_num[i * 9], self.filter_num[i * 9 + 1],
+                     self.filter_num[i * 9 + 2], self.filter_num[i * 9 + 3],
+                     self.filter_num[i * 9 + 4], self.filter_num[i * 9 + 5],
+                     self.filter_num[i * 9 + 6], self.k_size[i * 9 + 7], 2 if
+                     self.block_mask == 1 else 1, self.pool_type[i * 9 + 8]))
+        else:
+            repeat_num = self.block_num / self.downsample_num
+            num_minus = self.block_num % self.downsample_num
+            ### if block_num > downsample_num, add stride=1 block at last (block_num-downsample_num) layers
+            for i in range(self.downsample_num):
+                self.bottleneck_params_list.append(
+                    (self.filter_num[i * 9], self.filter_num[i * 9 + 1],
+                     self.filter_num[i * 9 + 2], self.filter_num[i * 9 + 3],
+                     self.filter_num[i * 9 + 4], self.filter_num[i * 9 + 5],
+                     self.filter_num[i * 9 + 6], self.k_size[i * 9 + 7], 2,
+                     self.pool_type[i * 9 + 8]))
+                ### if block_num / downsample_num > 1, add (block_num / downsample_num) times stride=1 block 
+                for k in range(repeat_num - 1):
+                    kk = k * self.downsample_num + i
+                    self.bottleneck_params_list.append(
+                        (self.filter_num[kk * 9], self.filter_num[kk * 9 + 1],
+                         self.filter_num[kk * 9 + 2],
+                         self.filter_num[kk * 9 + 3],
+                         self.filter_num[kk * 9 + 4],
+                         self.filter_num[kk * 9 + 5],
+                         self.filter_num[kk * 9 + 6], self.k_size[kk * 9 + 7],
+                         1, self.pool_type[kk * 9 + 8]))
+
+                if self.downsample_num - i <= num_minus:
+                    j = self.downsample_num * repeat_num + i
+                    self.bottleneck_params_list.append((
+                        self.filter_num[j * 9], self.filter_num[j * 9 + 1],
+                        self.filter_num[j * 9 + 2], self.filter_num[j * 9 + 3],
+                        self.filter_num[j * 9 + 4], self.filter_num[j * 9 + 5],
+                        self.filter_num[j * 9 + 6], self.k_size[j * 9 + 7], 1,
+                        self.pool_type[j * 9 + 8]))
+
+            if self.downsample_num == 0 and self.block_num != 0:
+                for i in range(len(self.block_num)):
+                    self.bottleneck_params_list.append((
+                        self.filter_num[i * 9], self.filter_num[i * 9 + 1],
+                        self.filter_num[i * 9 + 2], self.filter_num[i * 9 + 3],
+                        self.filter_num[i * 9 + 4], self.filter_num[i * 9 + 5],
+                        self.filter_num[i * 9 + 6], self.k_size[i * 9 + 7], 1,
+                        self.pool_type[i * 9 + 8]))
+
+        def net_arch(input, return_mid_layer=False, return_block=[]):
+            assert isinstance(return_block,
+                              list), 'return_block must be a list.'
+            layer_count = 0
+            mid_layer = dict()
+            for i, layer_setting in enumerate(self.bottleneck_params_list):
+                filter_nums = layer_setting[0:7]
+                filter_size = layer_setting[7]
+                stride = layer_setting[8]
+                pool_type = 'avg' if layer_setting[9] == 0 else 'max'
+                if stride == 2:
+                    layer_count += 1
+                if (layer_count - 1) in return_block:
+                    mid_layer[layer_count - 1] = input
+
+                input = self._inceptionA(
+                    input,
+                    layer_setting[0:7],
+                    filter_size=filter_size,
+                    stride=stride,
+                    pool_type=pool_type,
+                    name='inceptionA_{}'.format(i + 1))
+
+            if return_mid_layer:
+                return input, mid_layer
+            else:
+                return input
+
+        return net_arch
+
+    def _inceptionA(self,
+                    data,
+                    A_tokens,
+                    filter_size,
+                    stride,
+                    pool_type,
+                    name=None):
+        pool1 = fluid.layers.pool2d(
+            input=data,
+            pool_size=filter_size,
+            pool_padding='SAME',
+            pool_type=pool_type,
+            name=name + '_pool2d')
+        conv1 = conv_bn_layer(
+            input=pool1,
+            filter_size=1,
+            num_filters=A_tokens[0],
+            stride=stride,
+            act='relu',
+            name=name + '_conv1')
+
+        conv2 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=A_tokens[1],
+            stride=stride,
+            act='relu',
+            name=name + '_conv2')
+
+        conv3 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=A_tokens[2],
+            stride=1,
+            act='relu',
+            name=name + '_conv3_1')
+        conv3 = conv_bn_layer(
+            input=conv3,
+            filter_size=filter_size,
+            num_filters=A_tokens[3],
+            stride=stride,
+            act='relu',
+            name=name + '_conv3_2')
+
+        conv4 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=A_tokens[4],
+            stride=1,
+            act='relu',
+            name=name + '_conv4_1')
+        conv4 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=A_tokens[5],
+            stride=1,
+            act='relu',
+            name=name + '_conv4_2')
+        conv4 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=A_tokens[6],
+            stride=stride,
+            act='relu',
+            name=name + '_conv4_3')
+
+        concat = fluid.layers.concat(
+            [conv1, conv2, conv3, conv4], axis=1, name=name + '_concat')
+        return concat
+
+    def _inceptionB(self,
+                    data,
+                    B_tokens=[0] * 7,
+                    filter_size,
+                    stride,
+                    repeat,
+                    name=None):
+        pool1 = fluid.layers.pool2d(
+            input=data,
+            pool_size=filter_size,
+            pool_padding='SAME',
+            pool_type='avg',
+            name=name + '_inceptionB_pool2d')
+        conv1 = conv_bn_layer(
+            input=pool1,
+            filter_size=1,
+            num_filters=B_tokens[0],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionB_conv1')
+
+        conv2 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=B_tokens[1],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionB_conv2')
+
+        conv3 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=B_tokens[2],
+            stride=1,
+            act='relu',
+            name=name + '_inceptionB_conv3_1')
+        conv3 = conv_bn_layer(
+            input=conv3,
+            filter_size=filter_size,
+            num_filters=B_tokens[3],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionB_conv3_2')
+
+        conv4 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=B_tokens[4],
+            stride=1,
+            act='relu',
+            name=name + '_inceptionB_conv4_1')
+        conv4 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=B_tokens[5],
+            stride=1,
+            act='relu',
+            name=name + '_inceptionB_conv4_2')
+        conv4 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=B_tokens[6],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionB_conv4_3')
+        concat = fluid.layers.concat(
+            [conv1, conv2, conv3, conv4],
+            axis=1,
+            name=name + '_inceptionB_concat')
+        return concat
+
+    def _inceptionC(self,
+                    data,
+                    C_tokens=[0] * 9,
+                    filter_size,
+                    stride,
+                    repeat,
+                    name=None):
+        pool1 = fluid.layers.pool2d(
+            input=data,
+            pool_size=filter_size,
+            pool_padding='SAME',
+            pool_type='avg',
+            name=name + '_inceptionC_pool2d')
+        conv1 = conv_bn_layer(
+            input=pool1,
+            filter_size=1,
+            num_filters=C_tokens[0],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionC_conv1')
+
+        conv2 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=C_tokens[1],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionC_conv2')
+
+        conv3 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=C_tokens[2],
+            stride=1,
+            act='relu',
+            name=name + '_inceptionC_conv3_1')
+        conv3_1 = conv_bn_layer(
+            input=conv3,
+            filter_size=filter_size,
+            num_filters=C_tokens[3],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionC_conv3_2_1')
+        conv3_2 = conv_bn_layer(
+            input=conv3,
+            filter_size=filter_size,
+            num_filters=C_tokens[4],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionC_conv3_2_2')
+
+        conv4 = conv_bn_layer(
+            input=data,
+            filter_size=1,
+            num_filters=C_tokens[5],
+            stride=1,
+            act='relu',
+            name=name + '_inceptionC_conv4_1')
+        conv4 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=C_tokens[6],
+            stride=1,
+            act='relu',
+            name=name + '_inceptionC_conv4_2')
+        conv4_1 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=C_tokens[7],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionC_conv4_3_1')
+        conv4_2 = conv_bn_layer(
+            input=conv4,
+            filter_size=filter_size,
+            num_filters=C_tokens[8],
+            stride=stride,
+            act='relu',
+            name=name + '_inceptionC_conv4_3_2')
+
+        concat = fluid.layers.concat(
+            [conv1, conv2, conv3_1, conv3_2, conv4_1, conv4_2],
+            axis=1,
+            name=name + '_inceptionC_concat')
+        return concat

paddleslim/nas/search_space/mobilenet_block.py

@@ -40,12 +40,17 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
 
         if self.block_mask == None:
             # use input_size and output_size to compute self.downsample_num
-            self.downsample_num = compute_downsample_num(self.input_size, self.output_size)
+            self.downsample_num = compute_downsample_num(self.input_size,
+                                                         self.output_size)
         if self.block_num != None:
-            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(self.downsample_num, self.block_num)
+            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(
+                self.downsample_num, self.block_num)
 
         # self.filter_num means channel number
-        self.filter_num = np.array([3, 4, 8, 12, 16, 24, 32, 48, 64, 80, 96, 128, 144, 160, 192, 224, 256, 320, 384, 512]) # 20
+        self.filter_num = np.array([
+            3, 4, 8, 12, 16, 24, 32, 48, 64, 80, 96, 128, 144, 160, 192, 224,
+            256, 320, 384, 512
+        ])  # 20
         # self.k_size means kernel size
         self.k_size = np.array([3, 5])  #2
         # self.multiply means expansion_factor of each _inverted_residual_unit
@@ -58,22 +63,21 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
         if self.block_mask != None:
             return [0] * (len(self.block_mask) * 4)
         else:
-            return [0] * (self.block_num * 4) 
+            return [0] * (self.block_num * 4)
 
     def range_table(self):
         range_table_base = []
         if self.block_mask != None:
-            for i in range(len(self.block_mask)):
-                range_table_base.append(len(self.multiply))
-                range_table_base.append(len(self.filter_num))
-                range_table_base.append(len(self.repeat))
-                range_table_base.append(len(self.k_size))
+            range_table_length = len(self.block_mask)
         else:
-            for i in range(self.block_num):
-                range_table_base.append(len(self.multiply))
-                range_table_base.append(len(self.filter_num))
-                range_table_base.append(len(self.repeat))
-                range_table_base.append(len(self.k_size))
+            range_table_length = self.block_mum
+
+        for i in range(range_table_length):
+            range_table_base.append(len(self.multiply))
+            range_table_base.append(len(self.filter_num))
+            range_table_base.append(len(self.repeat))
+            range_table_base.append(len(self.k_size))
+
         return range_table_base
 
     def token2arch(self, tokens=None):
@@ -98,22 +102,40 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
             num_minus = self.block_num % self.downsample_num
             ### if block_num > downsample_num, add stride=1 block at last (block_num-downsample_num) layers
             for i in range(self.downsample_num):
-                self.bottleneck_params_list.append((self.mutiply[tokens[i * 4]], self.filter_num[tokens[i * 4 + 1]],
-                           self.repeat[tokens[i * 4 + 2]], 2, self.k_size[tokens[i * 4 + 3]]))
+                self.bottleneck_params_list.append(
+                    (self.multiply[tokens[i * 4]],
+                     self.filter_num[tokens[i * 4 + 1]],
+                     self.repeat[tokens[i * 4 + 2]], 2,
+                     self.k_size[tokens[i * 4 + 3]]))
 
                 ### if block_num / downsample_num > 1, add (block_num / downsample_num) times stride=1 block 
                 for k in range(repeat_num - 1):
                     kk = k * self.downsample_num + i
-                    self.bottleneck_params_list.append((self.mutiply[tokens[kk * 4]], self.filter_num[tokens[kk * 4 + 1]],
-                               self.repeat[tokens[kk * 4 + 2]], 1, self.k_size[tokens[kk * 4 + 3]]))
-                    
+                    self.bottleneck_params_list.append(
+                        (self.multiply[tokens[kk * 4]],
+                         self.filter_num[tokens[kk * 4 + 1]],
+                         self.repeat[tokens[kk * 4 + 2]], 1,
+                         self.k_size[tokens[kk * 4 + 3]]))
+
                 if self.downsample_num - i <= num_minus:
                     j = self.downsample_num * repeat_num + i
-                    self.bottleneck_params_list.append((self.mutiply[tokens[j * 4]], self.filter_num[tokens[j * 4 + 1]],
-                               self.repeat[tokens[j * 4 + 2]], 1, self.k_size[tokens[j * 4 + 3]]))
+                    self.bottleneck_params_list.append(
+                        (self.multiply[tokens[j * 4]],
+                         self.filter_num[tokens[j * 4 + 1]],
+                         self.repeat[tokens[j * 4 + 2]], 1,
+                         self.k_size[tokens[j * 4 + 3]]))
+
+            if self.downsample_num == 0 and self.block_num != 0:
+                for i in range(len(self.block_num)):
+                    self.bottleneck_params_list.append(
+                        (self.multiply[tokens[i * 4]],
+                         self.filter_num[tokens[i * 4 + 1]],
+                         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.'
+            assert isinstance(return_block,
+                              list), 'return_block must be a list.'
             # all padding is 'SAME' in the conv2d, can compute the actual padding automatic. 
             # bottleneck sequences
             in_c = int(32 * self.scale)
@@ -127,7 +149,7 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
                 if s == 2:
                     layer_count += 1
                 if (layer_count - 1) in return_block:
-                    mid_layer[layer_count] = depthwise_conv
+                    mid_layer[layer_count - 1] = depthwise_conv
 
                 input, depthwise_conv = self._invresi_blocks(
                     input=input,
@@ -137,9 +159,12 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
                     n=n,
                     s=s,
                     k=k,
-                    name='mobilenetv2_' + str(i+1))
+                    name='mobilenetv2_' + str(i + 1))
                 in_c = int(c * self.scale)
 
+            if layer_count in return_block:
+                mid_layer[layer_count] = depthwise_conv
+
             if return_mid_layer:
                 return input, mid_layer
             else:
@@ -259,18 +284,28 @@ class MobileNetV2BlockSpace(SearchSpaceBase):
         return last_residual_block, depthwise_output
 
 
-
 @SEARCHSPACE.register
 class MobileNetV1BlockSpace(SearchSpaceBase):
-    def __init__(self, input_size, output_size, block_num, block_mask=None, scale=1.0):
-        super(MobileNetV1BlockSpace, self).__init__(input_size, output_size, block_num, block_mask)
+    def __init__(self,
+                 input_size,
+                 output_size,
+                 block_num,
+                 block_mask=None,
+                 scale=1.0):
+        super(MobileNetV1BlockSpace, self).__init__(input_size, output_size,
+                                                    block_num, block_mask)
         # use input_size and output_size to compute self.downsample_num
-        self.downsample_num = compute_downsample_num(self.input_size, self.output_size)
+        self.downsample_num = compute_downsample_num(self.input_size,
+                                                     self.output_size)
         if self.block_num != None:
-            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(self.downsample_num, self.block_num)
+            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(
+                self.downsample_num, self.block_num)
 
         # self.filter_num means channel number
-        self.filter_num = np.array([3, 4, 8, 12, 16, 24, 32, 48, 64, 80, 96, 128, 144, 160, 192, 224, 256, 320, 384, 512, 576, 640, 768, 1024, 1048])
+        self.filter_num = np.array([
+            3, 4, 8, 12, 16, 24, 32, 48, 64, 80, 96, 128, 144, 160, 192, 224,
+            256, 320, 384, 512, 576, 640, 768, 1024, 1048
+        ])
         self.k_size = np.array([3, 5])
         self.scale = scale
 
@@ -302,37 +337,55 @@ class MobileNetV1BlockSpace(SearchSpaceBase):
         self.bottleneck_param_list = []
         if self.block_mask != None:
             for i in range(len(self.block_mask)):
-                self.bottleneck_params_list.append((self.filter_num[tokens[i * 3]], self.filter_num[tokens[i * 3 + 1]], 2 if self.block_mask[i] == 1 else 1, self.k_size[tokens[i * 3 + 2]]))
+                self.bottleneck_params_list.append(
+                    (self.filter_num[tokens[i * 3]],
+                     self.filter_num[tokens[i * 3 + 1]], 2
+                     if self.block_mask[i] == 1 else 1,
+                     self.k_size[tokens[i * 3 + 2]]))
         else:
             repeat_num = self.block_num / self.downsample_num
             num_minus = self.block_num % self.downsample_num
             for i in range(self.block_num):
                 ### if block_num > downsample_num, add stride=1 block at last (block_num-downsample_num) layers
-                self.bottleneck_params_list.append((self.filter_num[tokens[i * 3]], self.filter_num[tokens[i * 3 + 1]], 2, self.k_size[tokens[i * 3 + 2]]))
+                self.bottleneck_params_list.append(
+                    (self.filter_num[tokens[i * 3]],
+                     self.filter_num[tokens[i * 3 + 1]], 2,
+                     self.k_size[tokens[i * 3 + 2]]))
 
                 ### if block_num / downsample_num > 1, add (block_num / downsample_num) times stride=1 block 
                 for k in range(repeat_num - 1):
                     kk = k * self.downsample_num + i
-                    self.bottleneck_params_list.append((self.filter_num[tokens[kk * 3]], self.filter_num[tokens[kk * 3 + 1]],
-                               1, self.k_size[tokens[kk * 3 + 2]]))
-                    
+                    self.bottleneck_params_list.append(
+                        (self.filter_num[tokens[kk * 3]],
+                         self.filter_num[tokens[kk * 3 + 1]], 1,
+                         self.k_size[tokens[kk * 3 + 2]]))
+
                 if self.downsample_num - i <= num_minus:
                     j = self.downsample_num * repeat_num + i
-                    self.bottleneck_params_list.append((self.filter_num[tokens[j * 3]], self.filter_num[tokens[j * 3 + 1]],
-                               1, self.k_size[tokens[j * 3 + 2]]))
-
+                    self.bottleneck_params_list.append(
+                        (self.filter_num[tokens[j * 3]],
+                         self.filter_num[tokens[j * 3 + 1]], 1,
+                         self.k_size[tokens[j * 3 + 2]]))
+
+            if self.downsample_num == 0 and self.block_num != 0:
+                for i in range(len(self.block_num)):
+                    self.bottleneck_params_list.append(
+                        (self.filter_num[tokens[i * 3]],
+                         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.'
+            assert isinstance(return_block,
+                              list), 'return_block must be a list.'
             mid_layer = dict()
             layer_count = 0
-            
+
             for i, layer_setting in enumerate(self.bottleneck_params_list):
                 filter_num1, filter_num2, stride, kernel_size = layer_setting
                 if stride == 2:
                     layer_count += 1
                 if (layer_count - 1) in return_block:
-                    mid_layer[layer_count] = input
+                    mid_layer[layer_count - 1] = input
 
                 input = self._depthwise_separable(
                     input=input,
@@ -343,6 +396,7 @@ class MobileNetV1BlockSpace(SearchSpaceBase):
                     scale=self.scale,
                     kernel_size=kernel_size,
                     name='mobilenetv1_{}'.format(str(i + 1)))
+
             if return_mid_layer:
                 return input, mid_layer
             else:
@@ -375,4 +429,3 @@ class MobileNetV1BlockSpace(SearchSpaceBase):
             name=name + '_sep')
 
         return pointwise_conv
-

paddleslim/nas/search_space/mobilenetv1.py

@@ -28,11 +28,7 @@ __all__ = ["MobileNetV1Space"]
 
 @SEARCHSPACE.register
 class MobileNetV1Space(SearchSpaceBase):
-    def __init__(self,
-                 input_size,
-                 output_size,
-                 block_num,
-                 block_mask):
+    def __init__(self, input_size, output_size, block_num, block_mask):
         super(MobileNetV1Space, self).__init__(input_size, output_size,
                                                block_num, block_mask)
         # self.head_num means the channel of first convolution
@@ -62,7 +58,6 @@ class MobileNetV1Space(SearchSpaceBase):
         # self.repeat means repeat_num in forth downsample 
         self.repeat = np.array([1, 2, 3, 4, 5, 6])  #6
 
-
     def init_tokens(self):
         """
         The initial token.
@@ -136,8 +131,8 @@ class MobileNetV1Space(SearchSpaceBase):
              self.k_size[tokens[18]]))
         for i in range(self.repeat[tokens[19]]):
             self.bottleneck_param_list.append(
-                (self.filter_num7[tokens[20]],
-                 self.filter_num8[tokens[21]], 1, self.k_size[tokens[22]]))
+                (self.filter_num7[tokens[20]], self.filter_num8[tokens[21]], 1,
+                 self.k_size[tokens[22]]))
         # 512 1024 1024 1024
         self.bottleneck_param_list.append(
             (self.filter_num8[tokens[23]], self.filter_num9[tokens[24]], 2,
@@ -160,8 +155,11 @@ class MobileNetV1Space(SearchSpaceBase):
                         s = 1
                     self.bottleneck_params_list[i] = (f1, f2, s, ks)
 
-
-        def net_arch(input, scale=1.0, return_block=[], end_points=None, output_stride=None):
+        def net_arch(input,
+                     scale=1.0,
+                     return_block=[],
+                     end_points=None,
+                     output_stride=None):
             self.scale = scale
             _modify_bottle_params(output_stride)
 
@@ -186,6 +184,14 @@ class MobileNetV1Space(SearchSpaceBase):
             layer_count = 1
             for i, layer_setting in enumerate(bottleneck_param_list):
                 filter_num1, filter_num2, stride, kernel_size = layer_setting
+                if stride == 2:
+                    layer_count += 1
+                ### return_block and end_points means block num
+                if check_points((layer_count - 1), return_block):
+                    decode_ends[layer_count - 1] = input
+
+                if check_points((layer_count - 1), end_points):
+                    return input, decode_ends
                 input = self._depthwise_separable(
                     input=input,
                     num_filters1=filter_num1,
@@ -195,15 +201,16 @@ class MobileNetV1Space(SearchSpaceBase):
                     scale=self.scale,
                     kernel_size=kernel_size,
                     name='mobilenetv1_{}'.format(str(i + 1)))
-                layer_count += 1
-                ### return_block and end_points means block num
-                if check_points(layer_count, return_block):
-                    decode_ends[layer_count] = depthwise_output
 
-                if check_points(layer_count, end_points):
-                    return input, decode_ends
+            ### return_block and end_points means block num
+            if check_points(layer_count, end_points):
+                return input, decode_ends
 
-            input = fluid.layers.pool2d(input=input, pool_type='avg', global_pooling=True, name='mobilenetv1_last_pool')
+            input = fluid.layers.pool2d(
+                input=input,
+                pool_type='avg',
+                global_pooling=True,
+                name='mobilenetv1_last_pool')
 
             return input
 

paddleslim/nas/search_space/mobilenetv2.py

@@ -28,11 +28,7 @@ __all__ = ["MobileNetV2Space"]
 
 @SEARCHSPACE.register
 class MobileNetV2Space(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(MobileNetV2Space, self).__init__(input_size, output_size,
                                                block_num, block_mask)
         # self.head_num means the first convolution channel
@@ -54,7 +50,6 @@ class MobileNetV2Space(SearchSpaceBase):
         # self.repeat means repeat_num _inverted_residual_unit in each _invresi_blocks 
         self.repeat = np.array([1, 2, 3, 4, 5, 6])  #6
 
-
     def init_tokens(self):
         """
         The initial token.
@@ -140,8 +135,8 @@ class MobileNetV2Space(SearchSpaceBase):
                     self.bottleneck_params_list[i] = (t, c, n, s, ks)
 
         def net_arch(input,
-                     scale = 1.0,
-                     return_block = [],
+                     scale=1.0,
+                     return_block=[],
                      end_points=None,
                      output_stride=None):
             self.scale = scale
@@ -169,19 +164,21 @@ class MobileNetV2Space(SearchSpaceBase):
                 act='relu6',
                 name='mobilenetv2_conv1')
             layer_count = 1
-            if check_points(layer_count, return_block):
-                decode_ends[layer_count] = input
-
-            if check_points(layer_count, end_points):
-                return input, decode_ends
 
+            depthwise_output = None
             # bottleneck sequences
-            i = 1
             in_c = int(32 * self.scale)
             for layer_setting in self.bottleneck_params_list:
                 t, c, n, s, k = layer_setting
-                i += 1
-                #print(input)
+                if s == 2:
+                    layer_count += 1
+                ### return_block and end_points means block num
+                if check_points((layer_count - 1), return_block):
+                    decode_ends[layer_count - 1] = depthwise_output
+
+                if check_points((layer_count - 1), end_points):
+                    return input, decode_ends
+
                 input, depthwise_output = self._invresi_blocks(
                     input=input,
                     in_c=in_c,
@@ -192,14 +189,13 @@ class MobileNetV2Space(SearchSpaceBase):
                     k=k,
                     name='mobilenetv2_conv' + str(i))
                 in_c = int(c * self.scale)
-                layer_count += 1
 
-                ### return_block and end_points means block num
-                if check_points(layer_count, return_block):
-                    decode_ends[layer_count] = depthwise_output
+            ### return_block and end_points means block num
+            if check_points(layer_count, return_block):
+                decode_ends[layer_count] = depthwise_output
 
-                if check_points(layer_count, end_points):
-                    return input, decode_ends
+            if check_points(layer_count, end_points):
+                return input, decode_ends
 
             # last conv
             input = conv_bn_layer(

paddleslim/nas/search_space/resnet_block.py

@@ -30,12 +30,16 @@ __all__ = ["ResNetBlockSpace"]
 @SEARCHSPACE.register
 class ResNetBlockSpace(SearchSpaceBase):
     def __init__(input_size, output_size, block_num, block_mask):
-        super(ResNetSpace, self).__init__(input_size, output_size, block_num, block_mask)
+        super(ResNetSpace, self).__init__(input_size, output_size, block_num,
+                                          block_mask)
         # use input_size and output_size to compute self.downsample_num
-        self.downsample_num = compute_downsample_num(self.input_size, self.output_size)
+        self.downsample_num = compute_downsample_num(self.input_size,
+                                                     self.output_size)
         if self.block_num != None:
-            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(self.downsample_num, self.block_num)
-        self.filter_num = np.array([48, 64, 96, 128, 160, 192, 224, 256, 320, 384, 512, 640])
+            assert self.downsample_num <= self.block_num, 'downsample numeber must be LESS THAN OR EQUAL TO block_num, but NOW: downsample numeber is {}, block_num is {}'.format(
+                self.downsample_num, self.block_num)
+        self.filter_num = np.array(
+            [48, 64, 96, 128, 160, 192, 224, 256, 320, 384, 512, 640])
         ### TODO: use repeat to compute normal cell
         #self.repeat = [2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24]
         self.k_size = np.array([3, 5])
@@ -58,23 +62,37 @@ class ResNetBlockSpace(SearchSpaceBase):
         self.bottleneck_params_list = []
         if self.block_mask != None:
             for i in range(len(self.block_mask)):
-                self.bottleneck_params_list.append((self.num_filters[tokens[i * 2]], self.kernel_size[tokens[i * 2 + 1]],
-                          2 if self.block_mask[i] == 1 else 1]))
+                self.bottleneck_params_list.append(
+                    (self.num_filters[tokens[i * 2]],
+                     self.kernel_size[tokens[i * 2 + 1]], 2
+                     if self.block_mask[i] == 1 else 1))
         else:
             repeat_num = self.block_num / self.downsample_num
             num_minus = self.block_num % self.downsample_num
             for i in range(self.downsample_num):
-                self.bottleneck_params_list.append(self.num_filters[tokens[i * 2]], self.kernel_size[tokens[i * 2 + 1]],
-                          2)
+                self.bottleneck_params_list.append(
+                    self.num_filters[tokens[i * 2]],
+                    self.kernel_size[tokens[i * 2 + 1]], 2)
                 for k in range(repeat_num - 1):
                     kk = k * self.downsample_num + i
-                    self.bottleneck_params_list.append(self.num_filters[tokens[kk * 2]], self.kernel_size[tokens[kk * 2 + 1]], 1)
+                    self.bottleneck_params_list.append(
+                        self.num_filters[tokens[kk * 2]],
+                        self.kernel_size[tokens[kk * 2 + 1]], 1)
                 if self.downsample_num - i <= num_minus:
                     j = self.downsample_num * repeat_num + i
-                    self.bottleneck_params_list.append(self.num_filters[tokens[j * 2]], self.kernel_size[tokens[j * 2 + 1]], 1)
+                    self.bottleneck_params_list.append(
+                        self.num_filters[tokens[j * 2]],
+                        self.kernel_size[tokens[j * 2 + 1]], 1)
+
+            if self.downsample_num == 0 and self.block_num != 0:
+                for i in range(len(self.block_num)):
+                    self.bottleneck_params_list.append(
+                        self.num_filters[tokens[i * 2]],
+                        self.kernel_size[tokens[i * 2 + 1]], 1)
 
         def net_arch(input, return_mid_layer=False, return_block=[]):
-            assert isinstance(return_block, list), 'return_block must be a list.'
+            assert isinstance(return_block,
+                              list), 'return_block must be a list.'
             layer_count = 0
             mid_layer = dict()
             for layer_setting in self.bottleneck_params_list:
@@ -82,9 +100,14 @@ class ResNetBlockSpace(SearchSpaceBase):
                 if stride == 2:
                     layer_count += 1
                 if (layer_count - 1) in return_block:
-                    mid_layer[layer_count] = input
+                    mid_layer[layer_count - 1] = input
 
-                input = self._bottleneck_block(input=input, num_filters=filter_num, kernel_size=k_size, stride=stride, name = 'resnet' + str(i + 1))
+                input = self._bottleneck_block(
+                    input=input,
+                    num_filters=filter_num,
+                    kernel_size=k_size,
+                    stride=stride,
+                    name='resnet' + str(i + 1))
 
             if return_mid_layer:
                 return input, mid_layer
@@ -96,11 +119,21 @@ class ResNetBlockSpace(SearchSpaceBase):
     def _shortcut(self, input, ch_out, stride, name=None):
         ch_in = input.shape[1]
         if ch_in != ch_out or stride != 1:
-            return self.conv_bn_layer(input, ch_out, 1, stride, name = name + '_shortcut')
+            return conv_bn_layer(
+                input=input,
+                filter_size=1,
+                num_filters=ch_out,
+                stride=stride,
+                name=name + '_shortcut')
         else:
             return input
 
-    def _bottleneck_block(self, input, num_filters, kernel_size, stride, name=None):
+    def _bottleneck_block(self,
+                          input,
+                          num_filters,
+                          kernel_size,
+                          stride,
+                          name=None):
         conv0 = conv_bn_layer(
             input=input,
             num_filters=num_filters,
@@ -121,8 +154,7 @@ class ResNetBlockSpace(SearchSpaceBase):
             act=None,
             name=name + '_bottleneck_conv2')
 
-        short = self._shortcut(
-            input, num_filters * 4, stride, name=name)
+        short = self._shortcut(input, num_filters * 4, stride, name=name)
 
         return fluid.layers.elementwise_add(
             x=short, y=conv2, act='relu', name=name + '_bottleneck_add')