CBAM-ResNet tf实现

以下是CBAM-ResNet的TensorFlow实现：

import tensorflow as tf

def conv2d(inputs, filters, kernel_size, strides=1, padding='same', activation=None, name=None):
    return tf.layers.conv2d(inputs=inputs, filters=filters, kernel_size=kernel_size,
                            strides=strides, padding=padding, activation=activation, name=name)

def batch_norm(inputs, training, momentum=0.997, epsilon=1e-5, name=None):
    return tf.layers.batch_normalization(inputs=inputs, momentum=momentum, epsilon=epsilon,
                                         scale=True, training=training, name=name)

def relu(inputs, name=None):
    return tf.nn.relu(inputs, name=name)

def max_pool2d(inputs, pool_size, strides, padding='same', name=None):
    return tf.layers.max_pooling2d(inputs=inputs, pool_size=pool_size, strides=strides,
                                   padding=padding, name=name)

def avg_pool2d(inputs, pool_size, strides, padding='same', name=None):
    return tf.layers.average_pooling2d(inputs=inputs, pool_size=pool_size, strides=strides,
                                       padding=padding, name=name)

def cbam_block(inputs, reduction_ratio=0.5, name=None):
    with tf.variable_scope(name):
        # Channel attention
        channels = inputs.get_shape()[-1]
        avg_pool = tf.reduce_mean(inputs, axis=[1, 2], keepdims=True)
        assert avg_pool.get_shape()[1:] == (1, 1, channels)
        max_pool = tf.reduce_max(inputs, axis=[1, 2], keepdims=True)
        assert max_pool.get_shape()[1:] == (1, 1, channels)

        fc1 = conv2d(avg_pool, int(channels * reduction_ratio), kernel_size=1, name='fc1')
        assert fc1.get_shape()[1:] == (1, 1, int(channels * reduction_ratio))
        relu1 = relu(fc1, name='relu1')
        fc2 = conv2d(relu1, channels, kernel_size=1, name='fc2')
        assert fc2.get_shape()[1:] == (1, 1, channels)
        # channel attention的权重
        ch_attention = tf.sigmoid(fc2 + max_pool)

        # Spatial attention
        max_pool2d = tf.reduce_max(ch_attention, axis=-1, keepdims=True)
        assert max_pool2d.get_shape()[1:] == (1, 1, 1)
        avg_pool2d = tf.reduce_mean(ch_attention, axis=-1, keepdims=True)
        assert avg_pool2d.get_shape()[1:] == (1, 1, 1)
        # spatial attention的权重
        sp_attention = tf.sigmoid(max_pool2d + avg_pool2d)

        # 输出加权后的特征
        output = inputs * ch_attention * sp_attention
        return output

def cbam_resnet_block(inputs, filters, strides, training, projection_shortcut, reduction_ratio=0.5, name=None):
    with tf.variable_scope(name):
        shortcut = inputs
        if projection_shortcut is not None:
            shortcut = projection_shortcut(inputs)

        inputs = conv2d(inputs, filters, kernel_size=1, strides=1, name='conv1')
        inputs = batch_norm(inputs, training=training, name='bn1')
        inputs = relu(inputs, name='relu1')

        inputs = conv2d(inputs, filters, kernel_size=3, strides=strides, name='conv2')
        inputs = batch_norm(inputs, training=training, name='bn2')
        inputs = relu(inputs, name='relu2')

        inputs = cbam_block(inputs, reduction_ratio=reduction_ratio, name='cbam_block')

        inputs += shortcut
        inputs = relu(inputs, name='relu_output')
        return inputs

def cbam_resnet(inputs, num_blocks, filters, training, reduction_ratio=0.5, name=None):
    with tf.variable_scope(name):
        # 第一层
        inputs = conv2d(inputs, filters[0], kernel_size=7, strides=2, name='conv1')
        inputs = batch_norm(inputs, training=training, name='bn1')
        inputs = relu(inputs, name='relu1')
        inputs = max_pool2d(inputs, pool_size=3, strides=2, name='max_pool1')

        # resnet blocks
        for i in range(num_blocks):
            filters_block = filters[i+1]
            strides = 1
            if i == 0:
                strides = 2
            inputs = cbam_resnet_block(inputs, filters_block, strides, training,
                                        projection_shortcut=cbam_projection_shortcut if strides > 1 else None,
                                        reduction_ratio=reduction_ratio, name='block{}'.format(i+1))

        inputs = batch_norm(inputs, training=training, name='bn_last')
        inputs = relu(inputs, name='relu_last')
        inputs = avg_pool2d(inputs, pool_size=inputs.get_shape()[1:-1], strides=1, name='avg_pool')
        inputs = tf.layers.flatten(inputs, name='flatten')
        return inputs


def cbam_projection_shortcut(inputs, filters, strides, training, reduction_ratio=0.5, name=None):
    with tf.variable_scope(name):
        shortcut = conv2d(inputs, filters, kernel_size=1, strides=strides, name='conv_projection')
        shortcut = batch_norm(shortcut, training=training, name='bn_projection')
        shortcut = cbam_block(shortcut, reduction_ratio=reduction_ratio, name='cbam_projection')
        return shortcut

其中，`cbam_block` 和 `cbam_resnet_block` 实现CBAM模块和CBAM-ResNet模块，`cbam_resnet` 实现CBAM-ResNet网络的整体架构，`cbam_projection_shortcut` 实现CBAM-ResNet中的投影shortcut。