如何处理含None维度的张量乘法并保留批次维度得到目标形状

你原来的操作会合并batch维度（即你提到的None维度）的所有数据做全局矩阵乘法，自然会丢失batch维度。要保留batch维度，你需要做批量样本级别的矩阵乘法，每个样本单独计算特征外积后保留batch轴，有两种常用实现方式：

• 方案1：手动扩充维度后批量矩阵乘法
  给A、B分别新增轴把二维向量转为三维的批量矩阵，再做matmul运算，代码如下：

res = tf.matmul(tf.expand_dims(A, axis=-1), tf.expand_dims(B, axis=1))

原理说明：tf.expand_dims(A, axis=-1)会把维度为(None, HWC)的A转为(None, HWC, 1)，tf.expand_dims(B, axis=1)会把B转为(None, 1, HWC)，TensorFlow的matmul默认保留最前面的批量维度，只对最后两个维度做矩阵乘法，运算后得到的维度就是(None, HWC, HWC)。

• 方案2：使用einsum语法直接定义维度运算逻辑
  用einsum可以更简洁地定义维度映射规则，不需要手动调整维度，代码如下：

res = tf.einsum('ij,ik->ijk', A, B)

原理说明：表达式中i对应batch维度，j对应A的HWC维度，k对应B的HWC维度，规则指定输出保留三个维度i/j/k，自动完成每个样本的外积运算，输出维度完全符合要求。

补充说明：你之前的写法tf.matmul(tf.transpose(A),B)会把整个A的维度从(None, HWC)转置为(HWC, None)，和B的(None, HWC)做矩阵乘法后得到(HWC, HWC)，本质是把所有batch样本的特征交叉结果累加了，自然会丢失batch维度，这是和你的需求不符的核心原因。

内容的提问来源于stack exchange，提问作者Alex Z