Tensorflow 中的张量运算

Question

如图所示，我有 2 个张量张量 A 和张量 B。我想从张量 B 中获取每个 1x16 向量，并沿 A 的轴 1 和 2 与张量 A 的相应批次（轴 0 是每个张量中的批次维度）执行点积。我如何在张量流中做到这一点。任何有关代码的帮助将不胜感激。谢谢

我尝试了 einsum，但很困惑。

预期的输出形状应该是 (3,2,28,28)。

Answer 1

你可以这样做：

import tensorflow as tf

A = tf.placeholder(tf.float32, [3, 28, 28, 16])
B = tf.placeholder(tf.float32, [3, 2, 16])
A_exp = A[:, tf.newaxis]
B_exp = B[:, :, tf.newaxis, tf.newaxis]
C = tf.reduce_sum(A_exp * B_exp, axis=-1)
print(C)
# Tensor("Sum:0", shape=(3, 2, 28, 28), dtype=float32)