tensorflow中densenet加入SE注意力机制代码

以下是在 TensorFlow 中实现 DenseNet 加入 SE 注意力机制的代码：

import tensorflow as tf

# 定义 SE 模块
def se_module(input_tensor, ratio=16):
    # 获取输入张量的通道数
    channels = input_tensor.get_shape()[-1]

    # 计算 squeeze 层输出的形状
    squeeze_output_shape = [1, 1, channels]

    # 定义 squeeze 层，将输入张量压缩成一个张量
    squeeze = tf.keras.layers.GlobalAveragePooling2D()(input_tensor)

    # 定义 excitation 层，增强输入张量中的有用特征
    excitation = tf.keras.layers.Dense(units=channels // ratio, activation='relu')(squeeze)
    excitation = tf.keras.layers.Dense(units=channels, activation='sigmoid')(excitation)
    excitation = tf.reshape(excitation, [-1, 1, 1, channels])

    # 返回加权后的张量
    return input_tensor * excitation

# 定义 DenseNet 模型
def densenet_se(input_shape=(224, 224, 3), num_classes=1000, dense_blocks=4, dense_layers=-1, growth_rate=32, dropout_rate=0.2, bottleneck=False, compression=1.0, se_ratio=16):
    # 输入层
    inputs = tf.keras.Input(shape=input_shape)

    # 首先进行一个卷积操作，将输入的图像转化为特征图
    x = tf.keras.layers.Conv2D(filters=2 * growth_rate, kernel_size=7, strides=2, padding='same', use_bias=False)(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.MaxPool2D(pool_size=3, strides=2, padding='same')(x)

    # 定义 DenseNet 模型的密集块和过渡块
    num_features = 2 * growth_rate
    for i in range(dense_blocks - 1):
        x, num_features = dense_block_se(x, num_features, num_layers=dense_layers, growth_rate=growth_rate, dropout_rate=dropout_rate, bottleneck=bottleneck, se_ratio=se_ratio)
        x = transition_layer(x, num_features=num_features, compression=compression, dropout_rate=dropout_rate)
        num_features = int(num_features * compression)

    # 最后一个密集块没有过渡块
    x, num_features = dense_block_se(x, num_features, num_layers=dense_layers, growth_rate=growth_rate, dropout_rate=dropout_rate, bottleneck=bottleneck, se_ratio=se_ratio)

    # 输出层
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(units=num_classes, activation='softmax')(x)

    # 创建模型
    model = tf.keras.models.Model(inputs, x)

    return model

# 定义 DenseNet 的密集块
def dense_block_se(input_tensor, num_features, num_layers, growth_rate, dropout_rate, bottleneck, se_ratio):
    # 定义一个列表，用于存储密集块中所有的卷积层输出的特征图
    features_list = [input_tensor]

    # 创建 num_layers 个卷积层
    for i in range(num_layers):
        x = bn_relu_conv(input_tensor, growth_rate, dropout_rate, bottleneck=bottleneck)

        # 将当前卷积层的输出特征图添加到特征图列表中
        features_list.append(x)

        # 将所有的特征图拼接在一起
        x = tf.keras.layers.Concatenate(axis=-1)(features_list)

        # 使用 SE 注意力机制增强特征图
        x = se_module(x, ratio=se_ratio)

        # 更新 num_features
        num_features += growth_rate

    # 返回最后的特征图和更新后的 num_features
    return x, num_features

# 定义 DenseNet 的过渡块
def transition_layer(input_tensor, num_features, compression, dropout_rate):
    # 计算压缩后的通道数
    num_features = int(num_features * compression)

    # 定义 Batch Normalization 层
    x = tf.keras.layers.BatchNormalization()(input_tensor)

    # 定义卷积层
    x = tf.keras.layers.Conv2D(filters=num_features, kernel_size=1, padding='same', use_bias=False)(x)

    # 添加 Dropout 层
    if dropout_rate:
        x = tf.keras.layers.Dropout(dropout_rate)(x)

    # 定义平均池化层
    x = tf.keras.layers.AveragePooling2D(pool_size=2, strides=2)(x)

    return x

# 定义 BN -> ReLU -> Conv 的卷积块
def bn_relu_conv(input_tensor, filters, dropout_rate=None, bottleneck=False):
    # Batch Normalization 层
    x = tf.keras.layers.BatchNormalization()(input_tensor)

    # ReLU 激活函数
    x = tf.keras.layers.Activation('relu')(x)

    # 瓶颈卷积
    if bottleneck:
        x = tf.keras.layers.Conv2D(filters=filters * 4, kernel_size=1, use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation('relu')(x)
        x = tf.keras.layers.Conv2D(filters=filters, kernel_size=3, padding='same', use_bias=False)(x)
    # 普通卷积
    else:
        x = tf.keras.layers.Conv2D(filters=filters, kernel_size=3, padding='same', use_bias=False)(x)

    # Dropout 层
    if dropout_rate:
        x = tf.keras.layers.Dropout(dropout_rate)(x)

    return x

在上面的代码中，我们首先定义了一个 `se_module` 函数，用于定义 SE 注意力机制模块。在该函数中，我们首先获取输入特征图的通道数，然后使用 Global Average Pooling 层将输入特征图压缩成一个张量。接着，我们定义一个 squeeze 层和一个 excitation 层，其中 squeeze 层用于压缩通道数，excitation 层用于增强输入特征图中的有用特征。最后，我们将 squeeze 层和 excitation 层相乘，得到加权后的特征图。

接着，在 `densenet_se` 函数中，我们定义了 DenseNet 模型。首先，我们创建了一个输入层，并将输入的图像进行卷积操作，将其转化为特征图。然后，我们定义了 DenseNet 模型的密集块和过渡块，并使用 SE 注意力机制增强特征图。最后，我们添加了一个输出层，用于输出预测结果。

在 `dense_block_se` 函数中，我们定义了 DenseNet 的密集块。在该函数中，我们创建了 num_layers 个卷积层，并将所有的特征图拼接在一起。接着，我们使用 SE 注意力机制增强特征图，并返回最后的特征图和更新后的 num_features。

最后，在 `transition_layer` 函数中，我们定义了 DenseNet 的过渡块。在该函数中，我们首先计算压缩后的通道数，然后定义了 Batch Normalization 层、卷积层、Dropout 层和平均池化层，用于降低特征图的尺寸和通道数。