在Keras中实现简单的注意力机制

Question

我想实现一种简单的注意力机制来整合CNN模型的结果。

具体来说，我输入的每个示例都是一系列图像，因此每个示例的形状均为[None, img_width, img_height, n_channels]。

使用TimeDistributed包装器，可以应用CNN，以便获得形状为[None, hidden_state_size]的输出。

我想将CNN应用于序列中的每个图像，然后计算形状为[None]的注意力向量。为此，我通过带有单个输出单元的TimeDistributed Dense网络运行TimeDistributed CNN的输出，并计算Sequence上的softmax。

然后，将注意力向量乘以TimeDistributed CNN的输出，并对所有内容进行求和，以便最终得到形状为[hidden_state_size]的张量。

结果代码如下：

import tensorflow.keras as keras
import tensorflow.keras.layers as ll

inputs = ll.Input([None, 28, 28, 3])
x = inputs
x = ll.TimeDistributed(ll.Flatten())(x)
attention = ll.TimeDistributed(ll.Dense(1))(x)
attention = ll.Flatten()(attention)
attention = ll.Softmax()(attention)
outputs = ll.dot([x, attention], axes=[-2, -1])

model = keras.models.Model(inputs, outputs)

此模型的尺寸似乎已经检查出来，但这会满足我的要求吗？还是我在某个地方犯了错误？