pytorch扩展tensor的一个维度或多个维度

在使用神经网络的过程中，经常会用到把一个tensor扩展一个或多个维度的情况，然后把扩展后的维度用来广播。

• 扩展一个维度的情况，使用unsqueeze()函数，缩小一个维度使用squeeze()函数。参数是dim，是一个int。也就是说只能扩展或缩小一个维度。

t = torch.tensor([[1,2,3]])
t = t.unsqueeze(0)
#t.shape
#[1, 1, 2]

• 如果想扩展多个维度怎么办，答案是用tensor的切片来做，扩展一个维度也可以用切片来做。

# 把一个[2,3]的tensor扩展到[2,1,1,3]的tensor
>>> a = torch.tensor([[1,2,3],[4,5,6]])
>>> a
tensor([[1, 2, 3],
        [4, 5, 6]])
>>> b = a[:,None,None,:]
>>> b
tensor([[[[1, 2, 3]]],


        [[[4, 5, 6]]]])
>>> b.shape
torch.Size([2, 1, 1, 3])
>>> 

• 使用切片来扩展一个维度