PyTorch mean转TensorFlow tf.reduce_mean 结果不一致问题排查

问题本质是两个框架的特征张量维度布局天生存在差异，加上导入TensorFlow的张量没有按语义对齐维度，导致规约计算的对象从根上就错了：

• PyTorch的特征张量采用NCHW布局，也就是[批大小, 通道数, 高, 宽]，执行a.mean(dim=(2,3), keepdim=True)的逻辑，是每个通道单独计算自身覆盖的32*32个像素的算术平均值，输出每个通道对应一个独立标量，最终形状为[batch,27,1,1]。
• TensorFlow的同语义特征张量采用NHWC布局，也就是[批大小, 高, 宽, 通道数]，通道维度位于所有空间维度之后。你直觉上选择axis=(1,2)对应高、宽维度的均值计算，但数值出现巨大差异的核心原因是：你提供给TensorFlow的张量根本不是真正的NHWC语义布局——大概率是直接把PyTorch导出的NCHW格式连续内存数组，强行reshape成了[B,32,32,27]的形状，没有完成维度转置。
  这种情况下张量内部的元素位置完全错位，对axis=(1,2)求均值时，计算的是乱序混合后的跨通道、跨空间元素的平均值，根本不是预期的每个通道的独立空间均值，结果自然和PyTorch端差异巨大。
  举个最简单的对比例子：1个批次、2个通道、高2宽2的NCHW张量，按内存顺序排列的元素是[c0_h0_w0, c0_h0_w1, c0_h1_w0, c0_h1_w1, c1_h0_w0, c1_h0_w1, c1_h1_w0, c1_h1_w1]。如果直接把这8个元素拉平reshape成NHWC要求的[1,2,2,2]形状，得到的矩阵元素位置和真正转置后的NHWC张量完全错位，计算出的均值自然天差地别。

• 首先对齐TensorFlow张量的语义：如果张量是从PyTorch侧导出的，必须先转换为正确的NHWC布局，对应操作是a_tf = torch_a.permute(0, 2, 3, 1).numpy()再转换为TensorFlow张量，绝对不要直接对连续内存数组强行reshape成目标形状。
• 然后对正确布局的NHWC张量执行空间均值规约：运行tf.reduce_mean(a_tf, axis=(1,2), keepdims=True)，这一步计算出的均值数值和PyTorch端的结果是完全一致的，只是输出形状为[B,1,1,27]——通道维度被放在了最后一位。
• 如果需要和PyTorch端的输出形状完全匹配（即[B,27,1,1]），只需要把通道维度从最后一位移动到第二位即可，对应操作是result_tf = tf.transpose(reduced_result, perm=[0, 3, 1, 2])。完成以上步骤后，TensorFlow端的输出数值、形状都会和PyTorch端完全等效。

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