TensorFlow中三阶张量矩阵乘法问题求助

嘿，我完全懂你的困扰——想用三阶张量做对应维度的矩阵乘法，结果tf.matmul给出的形状完全不是你想要的，对吧？

先给你解释下为什么会这样：tf.matmul对高维张量的处理逻辑是，把除最后两个维度外的所有维度视为“批量（batch）维度”，而且要求两个输入的批量维度必须完全匹配。你的输入inputs是[None, n_type, n_features]，权重weight是[n_type, n_features, n_neurons]，这里inputs的批量维度是[None, n_type]，而weight的批量维度是[n_type]，两者不匹配，所以matmul错误地把这两个n_type维度当作匹配的批量维度运算，最终得到了[n_type, n_type, n_neurons]的奇怪结果。

下面给你几个可行的解决方案，按推荐程度排序：

1. 用tf.einsum（最直观，强烈推荐） #

einsum通过爱因斯坦求和符号直接定义维度间的运算关系，完美匹配你的需求：

Hidden1 = tf.einsum('b t f, t f n -> b t n', inputs, weight)

这里的符号含义很清晰：b代表你的batch维度（None），t是n_type，f是n_features，n是n_neurons。这个表达式就是告诉TensorFlow：对每个batch里的每个n_type，把对应n_features维度的向量和权重里同n_type的矩阵做乘法，最终输出[None, n_type, n_neurons]的张量，完全符合你的预期。

2. 重塑（Reshape）为二维张量做matmul，再恢复形状 #

如果你更习惯用matmul，可以先把张量拆成二维做运算，再还原维度：

# 将inputs从 [None, n_type, n_features] 重塑为 [None * n_type, n_features]
inputs_2d = tf.reshape(inputs, (-1, n_features))
# 将weight从 [n_type, n_features, n_neurons] 重塑为 [n_type * n_features, n_neurons]
weight_2d = tf.reshape(weight, (-1, n_neurons))
# 执行二维矩阵乘法
hidden1_2d = tf.matmul(inputs_2d, weight_2d)
# 把结果重塑回目标形状 [None, n_type, n_neurons]
Hidden1 = tf.reshape(hidden1_2d, (-1, n_type, n_neurons))

这个思路是把每个batch里的n_type个样本平铺成一个大的batch，和对应的权重矩阵做乘法后，再把结果拆分回原来的维度结构。

3. 用tf.map_fn逐个处理batch元素（适合小批量，不推荐大数据） #

如果你的batch规模不大，也可以用map_fn遍历每个batch元素，单独做矩阵乘法：

def process_batch_element(batch_input):
    # batch_input形状是 [n_type, n_features]，和weight做对应维度的乘法
    return tf.matmul(batch_input, weight)

Hidden1 = tf.map_fn(process_batch_element, inputs)

不过这个方法在大数据量下效率会比较低，所以优先考虑前两种方案。

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