TensorFlow中批量向量与矩阵的乘法实现问题

这个问题我之前也碰到过，tf.matmul()确实只支持二维矩阵的乘法运算，直接把三维的批量张量喂进去肯定会报维度不匹配的错。下面给你几个实用的解决办法：

方法1：用tf.einsum（最直观灵活） #

tf.einsum可以通过爱因斯坦求和约定来指定张量的运算规则，语法非常清晰，完全适配这种批量场景：

import tensorflow as tf

# 假设N=3，C=2的示例张量
x = tf.random.normal((3, 2))
Z = tf.random.normal((3, 2, 2))

# 执行批量行向量×矩阵运算
result = tf.einsum('nc,nck->nk', x, Z)
print(result.shape)  # 输出 (3, 2)，符合预期

解释一下：'nc,nck->nk'里的n代表批量维度，c是向量的维度，k是矩阵的列维度。这个表达式的意思是，对每个批量样本，取对应的1×C向量x[n]和C×C矩阵Z[n]相乘，最终得到N×C的结果张量。

方法2：扩展维度后用tf.matmul #

我们可以先把x的维度扩展成[N, 1, C]（给行向量加一个“行维度”），这样就能和[N, C, C]的Z做矩阵乘法，最后再把多余的维度去掉：

x_expanded = tf.expand_dims(x, axis=1)  # 维度变为 [N, 1, C]
temp_result = tf.matmul(x_expanded, Z)  # 乘法后维度是 [N, 1, C]
result = tf.squeeze(temp_result, axis=1)  # 去掉中间的1维度，得到 [N, C]

这个方法的核心是把批量里的每个行向量伪装成1×C的二维矩阵，让tf.matmul能识别并执行乘法，最后再还原成我们需要的形状。

方法3：用tf.linalg.matvec（专门的矩阵-向量乘法） #

TensorFlow的tf.linalg.matvec原生支持批量的矩阵乘列向量运算，不过因为你的需求是行向量乘矩阵，我们可以通过转置来适配：

# 行向量×矩阵 = (矩阵转置×列向量) 的转置
transposed_Z = tf.transpose(Z, perm=[0, 2, 1])  # Z转置后维度 [N, C, C]
x_col = tf.expand_dims(x, axis=-1)  # 把x转成列向量，维度 [N, C, 1]
temp_result = tf.linalg.matvec(transposed_Z, x_col)  # 结果维度 [N, C, 1]
result = tf.squeeze(temp_result, axis=-1)  # 还原成 [N, C]

不过相对前两种方法，这个步骤稍显繁琐，更推荐前两种。

内容的提问来源于stack exchange，提问作者v-i-s-h