是否可以使用 EinsumDense 而不是多个并行的 Dense 层？

Question

我有一个形状为 (None, 20, 250) 的输入，其中 20 是我的上下文窗口，250 是我的嵌入维度。

我想为上下文窗口中的每个元素应用一个 不同 密集的 250 -> 250。 以下代码可以正常工作并执行我想要的操作，但不使用 Einsum：

x = tf.unstack(x, axis=1)
x = [layers.Dense(250)(i) for i in x]
x = tf.stack(x, axis=1)

如果我对 Einsum 的理解正确，下面的代码应该做的完全一样：

einsum = EinsumDense('abc,cbd->abd', output_shape=[20, 250], bias_axes='bd')
x = einsum(x)

EinsumDense 的文档指出ab,bc->ac 相当于一个密集层，所以我认为这应该可行。 但是，这样做时我得到了截然不同的结果。

我在这里做错了什么？

Answer 1

试试这个！！！

einsum = EinsumDense('abc,cd->abd', output_shape=[20, 250], bias_axes='d')