cbam+attention 代码实战 完整的输入输出测试

CBAM（Convolutional Block Attention Module）是一种结合了通道注意力和空间注意力的注意力机制，用于增强卷积神经网络（CNN）的表示能力。下面是一个简单的示例，展示了如何在 Python 中实现 CBAM 模块，并进行完整的输入输出测试。

在示例中，我们将实现一个简化版本的 CBAM 模块，它包括通道注意力机制和空间注意力机制。

首先，我们创建 CBAM 模块的代码：

import tensorflow as tf

class ChannelAttention(tf.keras.layers.Layer):
    def __init__(self, ratio=8):
        super(ChannelAttention, self).__init__()
        self.ratio = ratio

    def build(self, input_shape):
        self.filters = input_shape[-1]
        self.shared_layer_one = tf.keras.layers.Dense(self.filters // self.ratio, activation='relu')
        self.shared_layer_two = tf.keras.layers.Dense(self.filters, activation='sigmoid')

    def call(self, inputs):
        avg_pool = tf.reduce_mean(inputs, axis=[1, 2], keepdims=True)
        avg_pool = self.shared_layer_one(avg_pool)
        channel_attention = self.shared_layer_two(avg_pool)
        return inputs * channel_attention

class SpatialAttention(tf.keras.layers.Layer):
    def __init__(self):
        super(SpatialAttention, self).__init__()
        self.conv = tf.keras.layers.Conv2D(filters=1, kernel_size=3, padding='same', activation='sigmoid')

    def call(self, inputs):
        max_pool = tf.reduce_max(inputs, axis=-1, keepdims=True)
        spatial_attention = self.conv(max_pool)
        return inputs * spatial_attention

class CBAM(tf.keras.layers.Layer):
    def __init__(self, ratio=8):
        super(CBAM, self).__init__()
        self.ca = ChannelAttention(ratio)
        self.sa = SpatialAttention()

    def call(self, inputs):
        x = self.ca(inputs)
        x = self.sa(x)
        return x

接下来，我们使用一个简单的模型来测试 CBAM 模块的效果：

# 创建一个简单的测试模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='relu', input_shape=(224, 224, 3)),
    CBAM(),  # 加入CBAM模块
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 生成随机输入数据进行测试
import numpy as np
X_train = np.random.rand(100, 224, 224, 3)
y_train = np.random.randint(0, 10, size=(100,))

# 训练模型
model.fit(X_train, y_train, epochs=5, batch_size=32, validation_split=0.2)

# 使用训练好的模型进行预测
X_test = np.random.rand(10, 224, 224, 3)
predictions = model.predict(X_test)
print(predictions)

在这个示例中，我们创建了一个简单的 CNN 模型，并在其中加入了 CBAM 模块。然后，我们使用随机生成的输入数据进行了训练和预测。实际应用中，你可能需要使用更大规模的数据集和更复杂的模型来评估 CBAM 模块的效果。