在使用神经网络的过程中,经常会用到把一个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])
>>> 
  • 使用切片来扩展一个维度
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐