在PyTorch中调换批次轴位置是否会对模型运行性能产生影响？

你观察到的逐元素算术运算保留输入步幅布局的结论是正确的，PyTorch的逐元素操作不会主动重排张量内存，只会继承输入张量的步幅属性。

两种写法存在内存局部性差异 #

PyTorch默认通过torch.zeros创建的张量为C风格（行优先）连续存储，维度越靠左，对应步幅越大，相邻索引的内存访问跨度越大。
两种写法的批次轴步幅差异如下：

• 第一种transpose(1,2)写法：初始创建的张量形状为(3, 批次大小, ADDITIONAL_DIM, x.shape[0]) + 剩余维度，批次对应初始第1轴，步幅为ADDITIONAL_DIM * x.shape[0] * 剩余维度元素总数，转置后批次轴移动到第2位，步幅保持不变。
• 第二种permute(1,2,0,...)写法：初始创建的张量形状为(批次大小, 3, ADDITIONAL_DIM, x.shape[0]) + 剩余维度，批次对应初始第0轴，步幅为3 * ADDITIONAL_DIM * x.shape[0] * 剩余维度元素总数，是第一种写法批次轴步幅的3倍，permute后批次轴移动到第2位，步幅仍保持这个更大的数值。

如果你的后续操作需要频繁沿批次轴访问数据，第一种写法的内存跨度更小，缓存命中率更高，运算性能更好。

是否需要关注这类性能影响 #

你可以根据实际使用场景判断：

• 如果张量尺寸较小、或者这部分运算不是整个程序的性能瓶颈，不需要额外关注，优先选择代码可读性更高、实现更简单的写法即可。
• 如果张量尺寸很大（批次过万、特征维度极高），且这部分张量是运算流程的高频操作对象，建议先做简单的性能测试验证差异：
  • 性能差异在10%以内的话不需要调整，开发效率优先
  • 性能差异超过30%的话，可以在维度调换后调用.contiguous()方法重新排布内存，后续操作的内存局部性会回到最优状态，仅需要承担一次一次性的内存拷贝开销。

如果你调整维度顺序仅仅是为了简化索引写法，也可以直接使用PyTorch内置的select方法，比如x.select(1, k)和x[:, k]效果完全一致，不需要调整维度顺序也能简化代码，避免布局调整带来的性能问题。

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