PyTorch中如何高效多次执行squeeze/unsqueeze操作？

假设我们有形状为[4,1,1,2]的张量：

import torch 
a = torch.tensor([[1,2,3,4],[5,6,7,8]], dtype=torch.float32)
a = a.reshape(4,1,1,2)
print(a.shape)  # 输出: torch.Size([4, 1, 1, 2])

想要将其压缩为[4,2]，但不想重复写squeeze(1).squeeze(1)，可以用以下几种简洁方式：

方案1：无参数调用squeeze() #

如果所有维度为1的轴都需要压缩，直接不带参数调用squeeze()即可，它会自动处理所有单维度：

a1 = a.squeeze()
print(a1.shape)  # 输出: torch.Size([4, 2])

方案2：指定多个维度压缩（精准控制） #

如果只想压缩特定的几个维度（比如存在其他需要保留的单维度），可以给squeeze的dim参数传入一个维度元组，一次性指定所有要压缩的轴：

a1 = torch.squeeze(a, dim=(1, 2))
# 或者用张量实例方法：a1 = a.squeeze(dim=(1, 2))
print(a1.shape)  # 输出: torch.Size([4, 2])

对应unsqueeze的批量操作 #

同样的逻辑也适用于unsqueeze，一次性给多个维度添加单维度：

# 从[4,2]变回[4,1,1,2]
a2 = torch.unsqueeze(a1, dim=(1, 2))
print(a2.shape)  # 输出: torch.Size([4, 1, 1, 2])

效率说明 #

以上所有方法和多次调用squeeze/unsqueeze的效率完全一致，因为这些操作都只是修改张量的形状元信息，不会涉及任何数据拷贝，属于O(1)的轻量操作。

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