trnsorflow下的CBAM模块加入三维卷积神经网络中

CBAM（Convolutional Block Attention Module）是一种用于增强卷积神经网络性能的注意力机制。它可以自适应地调整通道和空间特征的重要性，从而提高模型的准确性和泛化能力。

在TensorFlow中，可以通过自定义层来实现CBAM模块的加入。以下是一个简单的示例：

import tensorflow as tf

class ChannelAttention(tf.keras.layers.Layer):
    def __init__(self, reduction_ratio=8):
        super(ChannelAttention, self).__init__()
        self.reduction_ratio = reduction_ratio
    
    def build(self, input_shape):
        self.filters = input_shape[-1]
        self.avg_pool = tf.keras.layers.GlobalAveragePooling2D()
        self.max_pool = tf.keras.layers.GlobalMaxPooling2D()
        self.dense1 = tf.keras.layers.Dense(units=self.filters // self.reduction_ratio,
                                            activation='relu',
                                            kernel_initializer='he_normal',
                                            use_bias=True,
                                            bias_initializer='zeros')
        self.dense2 = tf.keras.layers.Dense(units=self.filters,
                                            activation='sigmoid',
                                            kernel_initializer='he_normal',
                                            use_bias=True,
                                            bias_initializer='zeros')
    
    def call(self, inputs):
        avg_pool = self.avg_pool(inputs)
        max_pool = self.max_pool(inputs)
        concat = tf.concat([avg_pool, max_pool], axis=-1)
        dense1 = self.dense1(concat)
        dense2 = self.dense2(dense1)
        reshape = tf.reshape(dense2, [-1, 1, 1, self.filters])
        scale = inputs * reshape
        return scale

class SpatialAttention(tf.keras.layers.Layer):
    def __init__(self):
        super(SpatialAttention, self).__init__()
    
    def build(self, input_shape):
        self.filters = input_shape[-1]
        self.conv1 = tf.keras.layers.Conv2D(filters=1,
                                            kernel_size=(3, 3),
                                            strides=(1, 1),
                                            padding='same',
                                            activation='sigmoid',
                                            kernel_initializer='he_normal',
                                            use_bias=False)
    
    def call(self, inputs):
        conv1 = self.conv1(inputs)
        scale = inputs * conv1
        return scale

class CBAM(tf.keras.layers.Layer):
    def __init__(self, reduction_ratio=8):
        super(CBAM, self).__init__()
        self.ChannelAttention = ChannelAttention(reduction_ratio=reduction_ratio)
        self.SpatialAttention = SpatialAttention()
    
    def build(self, input_shape):
        pass
    
    def call(self, inputs):
        ca = self.ChannelAttention(inputs)
        sa = self.SpatialAttention(ca)
        return sa

上述代码定义了一个CBAM层，包括通道注意力模块和空间注意力模块。通道注意力模块通过全局平均池化和全局最大池化对输入进行特征提取，并通过两个全连接层得到通道注意力系数，最终将其与输入相乘得到加强的特征。空间注意力模块通过一个卷积层得到空间注意力系数，将其与输入相乘得到加强的特征。

在三维卷积神经网络中，可以将上述代码中的Conv2D替换为Conv3D，然后将CBAM层添加到模型中即可：

import tensorflow as tf

class ChannelAttention(tf.keras.layers.Layer):
    def __init__(self, reduction_ratio=8):
        super(ChannelAttention, self).__init__()
        self.reduction_ratio = reduction_ratio
    
    def build(self, input_shape):
        self.filters = input_shape[-1]
        self.avg_pool = tf.keras.layers.GlobalAveragePooling3D()
        self.max_pool = tf.keras.layers.GlobalMaxPooling3D()
        self.dense1 = tf.keras.layers.Dense(units=self.filters // self.reduction_ratio,
                                            activation='relu',
                                            kernel_initializer='he_normal',
                                            use_bias=True,
                                            bias_initializer='zeros')
        self.dense2 = tf.keras.layers.Dense(units=self.filters,
                                            activation='sigmoid',
                                            kernel_initializer='he_normal',
                                            use_bias=True,
                                            bias_initializer='zeros')
    
    def call(self, inputs):
        avg_pool = self.avg_pool(inputs)
        max_pool = self.max_pool(inputs)
        concat = tf.concat([avg_pool, max_pool], axis=-1)
        dense1 = self.dense1(concat)
        dense2 = self.dense2(dense1)
        reshape = tf.reshape(dense2, [-1, 1, 1, 1, self.filters])
        scale = inputs * reshape
        return scale

class SpatialAttention(tf.keras.layers.Layer):
    def __init__(self):
        super(SpatialAttention, self).__init__()
    
    def build(self, input_shape):
        self.filters = input_shape[-1]
        self.conv1 = tf.keras.layers.Conv3D(filters=1,
                                            kernel_size=(3, 3, 3),
                                            strides=(1, 1, 1),
                                            padding='same',
                                            activation='sigmoid',
                                            kernel_initializer='he_normal',
                                            use_bias=False)
    
    def call(self, inputs):
        conv1 = self.conv1(inputs)
        scale = inputs * conv1
        return scale

class CBAM(tf.keras.layers.Layer):
    def __init__(self, reduction_ratio=8):
        super(CBAM, self).__init__()
        self.ChannelAttention = ChannelAttention(reduction_ratio=reduction_ratio)
        self.SpatialAttention = SpatialAttention()
    
    def build(self, input_shape):
        pass
    
    def call(self, inputs):
        ca = self.ChannelAttention(inputs)
        sa = self.SpatialAttention(ca)
        return sa

inputs = tf.keras.Input(shape=(32, 32, 32, 3))
x = tf.keras.layers.Conv3D(filters=64, kernel_size=(3, 3, 3), strides=(1, 1, 1), padding='same')(inputs)
x = CBAM()(x)
x = tf.keras.layers.Conv3D(filters=128, kernel_size=(3, 3, 3), strides=(2, 2, 2), padding='same')(x)
x = CBAM()(x)
x = tf.keras.layers.Conv3D(filters=256, kernel_size=(3, 3, 3), strides=(2, 2, 2), padding='same')(x)
x = CBAM()(x)
x = tf.keras.layers.Flatten()(x)
outputs = tf.keras.layers.Dense(units=10, activation='softmax')(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.summary()

上述代码定义了一个三维卷积神经网络，其中包括三个CBAM层。在模型训练时，可以像使用其他层一样进行编译和训练。