基于Tensorflow的Darknet53实现

官网：https://pjreddie.com/darknet/imagenet/

import tensorflow as tf
filter_initializer = tf.contrib.layers.xavier_initializer()


# 定义卷积操作
def conv_op(input_op, name, kh, kw, n_out, dh, dw, relu='leaky', normal=True, train=True):
    n_in = input_op.get_shape()[-1].value    # 获取输入图像的通道channel数
    with tf.variable_scope(name) as scope:         # 命名空间
        kernel = tf.get_variable("w",
                                 shape = [kh, kw, n_in, n_out],
                                 dtype = tf.float32,
                                 initializer = tf.truncated_normal_initializer(stddev=0.1, dtype=tf.float32)) #定义卷积核以及初始化方法

        biases = tf.get_variable('bases',
                                 shape=[n_out],
                                 dtype=tf.float32,
                                 initializer=tf.constant_initializer(0.1))

        conv = tf.nn.conv2d(input_op, kernel, strides=[1, dh, dw, 1], padding='SAME')
        conv = tf.nn.bias_add(conv, biases)

        if normal:
            conv = tf.layers.batch_normalization(inputs=conv, momentum=0.99, epsilon=1e-3, training=train,
                                          name=name+'_normal')
        if relu == 'leaky':
            conv = tf.nn.leaky_relu(conv, alpha = 0.1, name=name+relu)
        elif relu == 'lrn':
            tf.nn.lrn(conv, depth_radius=4, bias=1.0, alpha=0.01 / 9.0, beta=0.75, name=name+relu)
        elif relu == 'relu':
            conv = tf.nn.relu(conv, name=name+relu)
        else:
            pass
        return conv


#定义池化层
def mpool_op(input_op, name, kh, kw, dh, dw):
    pool = tf.nn.max_pool(input_op, ksize=[1, kh, kw, 1], strides=[1, dh, dw, 1], padding='SAME', name= name)
    return pool


#定义池化层
def avgpool_op(input_op, name, kh, kw, dh, dw, padding='SAME'):
    pool = tf.nn.avg_pool(input_op, ksize=[1, kh, kw, 1], strides=[1, dh, dw, 1], padding=padding, name=name)
    return pool


def unmaxpool_op(input_op, name, k):
    shape = input_op.get_shape()
    unpool = tf.image.resize_nearest_neighbor(input_op, [k * shape[1], k * shape[2]], name=name)
    return unpool


def side_branch(inputs, filters, upsample_factor):
    # 反置卷积
    kernel_size = [2 * upsample_factor, 2 * upsample_factor]
    outputs = tf.layers.conv2d_transpose(inputs,
                                         filters,
                                         kernel_size,
                                         strides=(upsample_factor, upsample_factor),
                                         padding='same',
                                         activation=None,
                                         use_bias=True,
                                         kernel_initializer=filter_initializer,
                                         kernel_regularizer=None)
    return outputs

#定义normlization层
def batch_normalization_layer(input_layer, name = None, training = True, norm_decay = 0.99, norm_epsilon = 1e-3):
    '''
    Introduction
    ------------
        对卷积层提取的feature map使用batch normalization
    Parameters
    ----------
        input_layer: 输入的四维tensor
        name: batchnorm层的名字
        trainging: 是否为训练过程
        norm_decay: 在预测时计算moving average时的衰减率
        norm_epsilon: 方差加上极小的数，防止除以0的情况
    Returns
    -------
        bn_layer: batch normalization处理之后的feature map
    '''
    bn_layer = tf.layers.batch_normalization(inputs=input_layer,
        momentum=norm_decay, epsilon = norm_epsilon, center=True,
        scale=True, training=training, name=name)
    return tf.nn.leaky_relu(bn_layer, alpha = 0.1)


#定义全连接操作
def fc_op(input_op, name, n_out):
    dim = input_op.get_shape()[-1].value
    with tf.variable_scope(name) as scope:
        weights = tf.get_variable('weights',
                                  shape=[dim, n_out],
                                  dtype=tf.float32,
                                  initializer=tf.truncated_normal_initializer(stddev=0.1, dtype=tf.float32))
        biases = tf.get_variable('bases',
                                 shape=[n_out],
                                 dtype=tf.float32,
                                 initializer=tf.constant_initializer(0.1))

        fc = tf.nn.leaky_relu(tf.matmul(input_op, weights) + biases, alpha=0.1, name='fc')
        return fc


def params_usage():
    """统计参数数量"""
    total = 0
    for v in tf.trainable_variables():
        shape = v.get_shape()
        cnt = 1
        for dim in shape:
            cnt *= dim.value
        total += cnt
    return total


def darknet53(inputs, filters=4):
    ###darknet53（53是指卷积层）网络结构
    conv1 = conv_op(inputs, name='conv1', kh=3, kw=3, n_out=filters, dh=1, dw=1, relu='leaky')

    # Downsample
    conv2 = conv_op(conv1, name='conv2', kh=3, kw=3, n_out=filters*2, dh=2, dw=2, relu='leaky')
    conv3 = conv_op(conv2, name='conv3', kh=1, kw=1, n_out=filters, dh=1, dw=1, relu='leaky')
    conv4 = conv_op(conv3, name='conv4', kh=3, kw=3, n_out=filters*2, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer5 = tf.concat([conv4, conv2], axis=-1, name='reslayer5')
    # Downsample
    conv6 = conv_op(reslayer5, name='conv6', kh=3, kw=3, n_out=filters*4, dh=2, dw=2, relu='leaky')
    conv7 = conv_op(conv6, name='conv7', kh=1, kw=1, n_out=filters*2, dh=1, dw=1, relu='leaky')
    conv8 = conv_op(conv7, name='conv8', kh=3, kw=3, n_out=filters*4, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer9 = tf.concat([conv8, conv6], axis=-1, name='reslayer9')
    conv10 = conv_op(reslayer9, name='conv10', kh=1, kw=1, n_out=filters*2, dh=1, dw=1, relu='leaky')
    conv11 = conv_op(conv10, name='conv11', kh=3, kw=3, n_out=filters*4, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer12 = tf.concat([conv11, reslayer9], axis=-1, name='reslayer12')

    # Downsample
    conv13 = conv_op(reslayer12, name='conv13', kh=3, kw=3, n_out=filters*8, dh=2, dw=2, relu='leaky')
    conv14 = conv_op(conv13, name='conv14', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv15 = conv_op(conv14, name='conv15', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer16 = tf.concat([conv15, conv13], axis=-1, name='reslayer16')
    conv17 = conv_op(reslayer16, name='conv17', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv18 = conv_op(conv17, name='conv18', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer19 = tf.concat([conv18, reslayer16], axis=-1, name='reslayer19')
    conv20 = conv_op(reslayer19, name='conv20', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv21 = conv_op(conv20, name='conv21', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer22 = tf.concat([conv21, reslayer19], axis=-1, name='reslayer22')
    conv23 = conv_op(reslayer22, name='conv23', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv24 = conv_op(conv23, name='conv24', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer25 = tf.concat([conv24, reslayer22], axis=-1, name='reslayer25')
    conv26 = conv_op(reslayer25, name='conv26', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv27 = conv_op(conv26, name='conv27', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer28 = tf.concat([conv27, reslayer25], axis=-1, name='reslayer28')
    conv29 = conv_op(reslayer28, name='conv29', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv30 = conv_op(conv29, name='conv30', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer31 = tf.concat([conv30, reslayer28], axis=-1, name='reslayer31')
    conv33 = conv_op(reslayer31, name='conv33', kh=1, kw=1, n_out=filters*4, dh=1, dw=1, relu='leaky')
    conv34 = conv_op(conv33, name='conv34', kh=3, kw=3, n_out=filters*8, dh=1, dw=1, relu='leaky')

    # Downsample
    conv35 = conv_op(conv34, name='conv35', kh=3, kw=3, n_out=filters*16, dh=2, dw=2, relu='leaky')
    conv36 = conv_op(conv35, name='conv36', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv37 = conv_op(conv36, name='conv37', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer38 = tf.concat([conv37, conv35], axis=-1, name='reslayer38')
    conv39 = conv_op(reslayer38, name='conv39', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv40 = conv_op(conv39, name='conv40', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer41 = tf.concat([conv40, reslayer38], axis=-1, name='reslayer41')
    conv42 = conv_op(reslayer41, name='conv42', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv43 = conv_op(conv42, name='conv43', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer44 = tf.concat([conv43, reslayer41], axis=-1, name='reslayer44')
    conv45 = conv_op(reslayer44, name='conv45', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv46 = conv_op(conv45, name='conv46', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer47 = tf.concat([conv46, reslayer44], axis=-1, name='reslayer47')
    conv48 = conv_op(reslayer47, name='conv48', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv49 = conv_op(conv48, name='conv49', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer50 = tf.concat([conv49, reslayer47], axis=-1, name='reslayer50')
    conv51 = conv_op(reslayer50, name='conv51', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv52 = conv_op(conv51, name='conv52', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer53 = tf.concat([conv52, reslayer50], axis=-1, name='reslayer53')
    conv54 = conv_op(reslayer53, name='conv54', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv55 = conv_op(conv54, name='conv55', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer56 = tf.concat([conv55, reslayer53], axis=-1, name='reslayer56')
    conv57 = conv_op(reslayer56, name='conv57', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv58 = conv_op(conv57, name='conv58', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer59 = tf.concat([conv58, reslayer56], axis=-1, name='reslayer59')
    conv60 = conv_op(reslayer59, name='conv60', kh=1, kw=1, n_out=filters*8, dh=1, dw=1, relu='leaky')
    conv61 = conv_op(conv60, name='conv61', kh=3, kw=3, n_out=filters*16, dh=1, dw=1, relu='leaky')

    # Downsample
    conv62 = conv_op(conv61, name='conv62', kh=3, kw=3, n_out=filters*32, dh=2, dw=2, relu='leaky')
    conv63 = conv_op(conv62, name='conv63', kh=1, kw=1, n_out=filters*16, dh=1, dw=1, relu='leaky')
    conv64 = conv_op(conv63, name='conv64', kh=3, kw=3, n_out=filters*32, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer65 = tf.concat([conv64, conv62], axis=-1, name='reslayer65')
    conv66 = conv_op(reslayer65, name='conv66', kh=1, kw=1, n_out=filters*16, dh=1, dw=1, relu='leaky')
    conv67 = conv_op(conv66, name='conv67', kh=3, kw=3, n_out=filters*32, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer68 = tf.concat([conv67, reslayer65], axis=-1, name='reslayer68')
    conv69 = conv_op(reslayer68, name='conv69', kh=1, kw=1, n_out=filters*16, dh=1, dw=1, relu='leaky')
    conv70 = conv_op(conv69, name='conv70', kh=3, kw=3, n_out=filters*32, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer71 = tf.concat([conv70, reslayer68], axis=-1, name='reslayer71')
    conv72 = conv_op(reslayer71, name='conv72', kh=1, kw=1, n_out=filters*16, dh=1, dw=1, relu='leaky')
    conv73 = conv_op(conv72, name='conv73', kh=3, kw=3, n_out=filters*32, dh=1, dw=1, relu='leaky')
    # shortcut
    reslayer74 = tf.concat([conv73, reslayer71], axis=-1, name='reslayer74')

    ###以上为darknet53（53是指卷积层）网络结构
    return reslayer74