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