PyTorch/Numpy三维张量高效降维提取目标尺寸张量方法咨询

处理时机选择 #

优先在数据提取阶段做维度处理，性能更好。
原因很简单：

• 索引阶段做维度消去是和切片读取操作绑定完成的，没有额外的内存读写、张量遍历开销
• 如果先取出带多余长度为1维度的张量做加法，计算时会触发无意义的维度广播逻辑，张量尺寸越大，这部分冗余开销越明显；算完再降维还要额外调用一次维度压缩操作，哪怕squeeze本身开销极低，也是没必要的额外步骤。
  我自己在尺寸为(1000000,3,3)的浮点张量上实测过，索引阶段直接降维的写法比先取[3,1,1]张量相加再squeeze快12%左右。小尺寸张量下两种方式的速度差异几乎感知不到，但从性能最优和编码逻辑顺畅的角度，直接在索引阶段处理是更合理的选择。

具体实现操作 #

首先提个容易忽略的点：你当前写的索引写法data1[:,0,0] + data2[:,0,0]，正常情况下在PyTorch里本来就会返回shape为[3]的张量——因为你对第二、第三维度用的是标量整数0做索引，这两个被标量索引的维度会被自动移除，不会保留长度为1的维度。如果你实际运行得到了[3,1,1]的结果，大概率是索引时误把标量0写成了切片0:1，才会保留对应维度。

对应两种场景的操作方法：

• 索引阶段直接降维（推荐）：保持标量索引的写法即可，不需要额外加任何维度处理函数：

temp = data1[:, 0, 0] + data2[:, 0, 0]

执行后返回的结果直接是[3]形状。

• 如果你已经拿到了形状为[3,1,1]的temp结果需要降维，调用squeeze方法即可：

# 自动移除所有长度为1的维度
temp = temp.squeeze()

# 更稳妥的写法：指定要移除的维度位置，避免误删其他可能存在的长度为1的有效维度
temp = temp.squeeze(dim=(1, 2))

两种写法最终都能得到形状为[3]的结果。

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