如何在TensorFlow中为含2层编码/解码Density层的自编码器瓶颈层加注意力层？

整体架构说明 #

你的自编码器结构可分为三部分：编码器（2个Dense层） → 瓶颈注意力层 → 解码器（2个Dense层）。核心是通过注意力层对瓶颈处的压缩特征进行加权，让模型更关注关键特征，提升重构效果。

代码实现步骤 #

1. 导入依赖库 #

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Layer
from tensorflow.keras.models import Model

2. 实现瓶颈注意力层 #

这里针对一维瓶颈特征实现了简单的自注意力机制，你也可以根据需求替换为多头注意力：

class BottleneckAttention(Layer):
    def __init__(self, units):
        super().__init__()
        self.query = Dense(units, activation='relu')
        self.key = Dense(units, activation='relu')
        self.value = Dense(units, activation='relu')
        self.softmax = tf.keras.layers.Softmax(axis=-1)

    def call(self, inputs):
        # 将瓶颈特征映射到query、key、value空间
        q = self.query(inputs)
        k = self.key(inputs)
        v = self.value(inputs)

        # 计算注意力权重并加权特征
        attention_scores = tf.matmul(tf.expand_dims(q, 1), tf.expand_dims(k, 2))
        attention_weights = self.softmax(attention_scores)
        output = tf.matmul(attention_weights, tf.expand_dims(v, 1))
        
        # 压缩回原瓶颈维度
        return tf.squeeze(output, axis=1)

3. 构建完整自编码器 #

def build_attention_autoencoder(input_dim, bottleneck_dim):
    # 编码器
    input_layer = Input(shape=(input_dim,))
    enc_dense1 = Dense(256, activation='relu')(input_layer)
    enc_dense2 = Dense(bottleneck_dim, activation='relu')(enc_dense1)

    # 插入瓶颈注意力层
    attention_out = BottleneckAttention(bottleneck_dim)(enc_dense2)

    # 解码器
    dec_dense1 = Dense(256, activation='relu')(attention_out)
    dec_dense2 = Dense(input_dim, activation='sigmoid')(dec_dense1)  # 图像任务用sigmoid，数值任务可换relu

    # 组装模型
    return Model(inputs=input_layer, outputs=dec_dense2)

4. 测试模型 #

以MNIST数据集（输入维度784）为例：

input_dim = 784
bottleneck_dim = 64

autoencoder = build_attention_autoencoder(input_dim, bottleneck_dim)
autoencoder.compile(optimizer='adam', loss='mse')
autoencoder.summary()

可选优化方案 #

• 多头注意力替换：如果需要更强的特征捕捉能力，可改用MultiHeadAttention层，只需调整注意力层实现：

class MultiHeadBottleneckAttention(Layer):
    def __init__(self, num_heads, key_dim):
        super().__init__()
        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=key_dim)

    def call(self, inputs):
        # 扩展维度适配多头注意力的序列输入要求
        inputs_expanded = tf.expand_dims(inputs, axis=1)
        attention_out = self.mha(query=inputs_expanded, value=inputs_expanded, key=inputs_expanded)
        return tf.squeeze(attention_out, axis=1)

替换时只需将BottleneckAttention改为MultiHeadBottleneckAttention(bottleneck_dim//2, bottleneck_dim)这类参数即可。

• 激活函数与损失调整：根据任务类型（图像、数值序列等）调整解码器输出的激活函数（如relu、tanh）和损失函数（如binary_crossentropy、mae）。

内容的提问来源于stack exchange，提问作者Filip