MobileNetOps

import tensorflow as tf

"""
实现Mobile_Net_V3的各个模块。
"""
weight_decay = 1e-3

def hswish(x, name='hswish'):
    return tf.multiply(x, tf.nn.relu6(x+3) / 6, name=name)


def batch_norm(x, momentum=0.99, epsilon=1e-5, train=True, name='bn'):
    return tf.layers.batch_normalization(x, momentum=momentum, epsilon=epsilon, training=train, name=name)


def conv2d(input, output_dims, k_h, k_w, d_h, d_w, stddev=0.1, name='conv2d', use_bias=False):
    """
    :param input:  tensor 4-D  [N, H, W, C]
    :param output_dims:   输出图片的depth(等于滤波器的个数)
    :param k_h:   池化核的高
    :param k_w:    池化核的宽
    :param d_h:    高方向上的步幅
    :param d_w:    宽方向上的步幅
    :param stddev:
    :param name:
    :param use_bias:
    """
    with tf.variable_scope(name):
        kernel = tf.get_variable(
            'w', shape=[k_h, k_w, input.get_shape()[-1], output_dims],
            initializer=tf.truncated_normal_initializer(stddev=stddev),
            regularizer=tf.contrib.layers.l2_regularizer(weight_decay)  # 批归一化和dropout有冲突, 一般用正则项,防止过拟合
        )
        conv = tf.nn.conv2d(input, kernel, strides=[1, d_h, d_w, 1], padding='SAME')
        if use_bias:
            bias = tf.get_variable('b', shape=[output_dims], initializer=tf.zeros_initializer())
            conv = tf.nn.bias_add(conv, bias)
        return conv


def conv2d_block(input, output_dim, k, s, bn_train, name, activation=hswish):
    """
    实现 1、卷积 + 2、BN + 3、激活
    :param input:
    :param output_dim:
    :param k:
    :param s:
    :param bn_train:
    :param name:
    :param activation:
    :return:
    """
    with tf.variable_scope(name):
        net = conv2d(input, output_dim, k, k, s, s, name='conv2d', use_bias=False)
        net = batch_norm(net, train=bn_train, name='bn')
        net = activation(net)
        return net


def conv_11(input, output_dims, name, use_bias=False):
    """
    实现1*1 卷积
    :param input:
    :param output_dims:
    :param name:
    :param bias:
    :return:
    """
    with tf.variable_scope(name):
        return conv2d(input, output_dims, 1, 1, 1, 1, name=name, use_bias=use_bias)


def pointwise_block(input, output_dim, bn_train, name, use_bias=False):
    """
    实现 1、1*1卷积 + 2、BN + 3、 激活
    :param input:
    :param output_dim:
    :param bn_train:
    :param name:
    :param use_bias:
    :return:
    """
    with tf.variable_scope(name):
        net = conv_11(input, output_dim, name, use_bias=use_bias)
        net = batch_norm(net, train=bn_train, name='bn')
        net = hswish(net)
        return net


def depthwise_conv(input, k_h=5, k_w=5, channel_multiplier=1, strides=[1,1,1,1],
                   padding='SAME', stddev=0.1, name='dwise_conv', use_bias=False):
    """
    实现深度智能卷积(depthwise-Conv)
    :param input:
    :param k_h:
    :param k_w:
    :param channel_multiplier:
    :param strides:
    :param padding:
    :param stddev:
    :param name:
    :param use_biae:
    :return:
    """
    with tf.variable_scope(name):
        in_channel = input.get_shape().as_list()[-1]
        kernel = tf.get_variable(
            'w', shape=[k_h, k_w, in_channel, channel_multiplier],
            initializer=tf.truncated_normal_initializer(stddev=stddev),
            regularizer=tf.contrib.layers.l2_regularizer(weight_decay)
        )
        net = tf.nn.depthwise_conv2d(input, filter=kernel, strides=strides, padding=padding)
        if use_bias:
            bias = tf.get_variable('b', shape=[in_channel * channel_multiplier], initializer=tf.zeros_initializer())
            net = tf.nn.bias_add(net, bias)
        return net


def Se_block(inputs_tensor, name, reduction=8):
    """
    :param inputs_tensor:  4-D tensor  [N, H, W, C]
    :param reduction:
    """
    with tf.variable_scope(name):
        _, H, W, C = inputs_tensor.get_shape()
        # 1、挤压
        x_mean = tf.reduce_mean(inputs_tensor, axis=[1, 2], keep_dims=True)
        # [N, 1, 1, C]
        # 2、激励   [N, 1, 1, C/r]
        x = tf.layers.conv2d(x_mean, C // reduction, kernel_size=1, strides=1, padding='valid', activation=tf.nn.relu,name='se1')
        # [N, 1, 1, C]
        x = tf.layers.conv2d(x, C, kernel_size=1, strides=1, padding='valid', activation=tf.nn.sigmoid, name='se2')
        # 3、特征重标定。
        y = tf.multiply(inputs_tensor, x)
    return y


def res_block(input, expansion_ratio, output_dims, stride, bn_train, name, activation=hswish,use_se=True, se_reduction=8, use_bias=False, shortcut=True):
    """
    实现MobileV3的瓶颈结构
    :param input:
    :param expansion_ratio:
    :param output_dims:
    :param stride:
    :param bn_train:
    :param name:
    :param activation:
    :param use_se:
    :param se_reduction:
    :param use_bias:
    :param shortcut:
    """
    with tf.variable_scope(name):
        bottleneck_dims = round(expansion_ratio * input.get_shape().as_list()[-1])
        # 1、1*1卷积
        net = conv_11(input, bottleneck_dims, name='pw', use_bias=use_bias)
        net = batch_norm(net, train=bn_train, name='pw_bn')
        net = activation(net)

        # 2、深度智能卷积（当strides=2, 代表通过深度智能卷积进行下采样）
        net = depthwise_conv(net, strides=[1, stride, stride, 1], name='dw', use_bias=use_bias)
        net = batch_norm(net, train=bn_train, name='dw_bn')
        net = activation(net)

        # 添加注意力操作
        if use_se:
            net = Se_block(net, name='SE', reduction=se_reduction)
        # 3、 1*1卷积（线性的）
        net = conv_11(net, output_dims, name='pw_linear', use_bias=use_bias)
        net = batch_norm(net, train=bn_train, name='pw_linear_bn')

        # 只在stride==1时候，做shortcut。 当stride=2时候（下采样时候），不做shortcut的。
        if shortcut and stride ==1:
            in_dim = int(input.get_shape().as_list()[-1])
            if in_dim != output_dims:
                ins = conv_11(input, output_dims, name='ex_dim')
                net = ins + net
            else:
                net = input + net
        return net

def global_avg(x):
    with tf.name_scope('global_avg'):
        net = tf.layers.average_pooling2d(x, pool_size=x.get_shape()[1:-1], strides=1)
        return net


def flatten(x):
    with tf.name_scope('flatten'):
        shape = x.get_shape()
        flatten_shape = shape[1] * shape[2] * shape[3]
        x= tf.reshape(x, shape=[-1, flatten_shape])
        return x