替换LSTM为Transformer实现层级注意力文档分类时形状不兼容报错

问题原因 #

你当前的模型输出形状是(None, 15, 5)，而标签的形状是(None,5)，二者维度不匹配的核心原因是：
经过句子级Transformer块处理后，输出的张量形状为(batch_size, max_sentences=15, 隐藏层维度)，你没有对句子维度做聚合，直接在后面接了输出层Dense(class_number, activation='softmax')，Dense层默认只作用于最后一维，所以会给每个句子都输出一个分类结果，而你需要的是整个文档对应一个分类结果，所以需要先把句子维度聚合掉。

解决方法 #

在句子级Transformer的输出后，添加一个全局平均池化（或全局最大池化）层，把句子维度压缩，再接入后续的全连接层和输出层即可。修改后的上层模型代码如下：

# Upper level (produce a representation of each document):
L2_dense_units = 100

sentence_input = layers.Input(shape=(max_sentences, max_words), name='sentence_input')

sentence_encoder = tf.keras.layers.TimeDistributed(word_encoder, name='sentence_encoder')(sentence_input)

sentence_transformer = TransformerBlock(embed_dim=L1_dense_units, num_heads=num_heads, ff_dim=ff_dim,
                               dropout_rate=dropout_rate, name='sentence_transformer')(sentence_encoder)
# 新增全局平均池化层，聚合所有句子的表征得到文档级表征
doc_pool = layers.GlobalAveragePooling1D(name='doc_pooling')(sentence_transformer)
sentence_dense = layers.Dense(int(L2_dense_units), activation='relu', name='sentence_dense')(doc_pool)
sentence_out = layers.Dropout(dropout_rate)(sentence_dense)
preds = layers.Dense(class_number , activation='softmax', name='sentence_output')(sentence_out)

model = keras.Model(sentence_input, preds)
model.summary()

修改说明 #

• 新增GlobalAveragePooling1D层，将形状从(batch_size, 15, 100)压缩为(batch_size, 100)，消除了句子维度
• 去掉了Dense层外面的TimeDistributed包装，现在Dense层直接作用于文档级表征，输出形状为(batch_size, 5)，和标签形状完全匹配，即可正常训练。

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