PyTorch中squeeze()和unsqueeze()详解

pytorch中squeeze()和unsqueeze()作用

---

squeeze()

squeeze() 用于在张量的指定维度插入新的维度 (为1) 得到维度提升的张量。

unsqueeze()

unsqueeze() 用于移除指定或者所有维度大小为1的维度，从而得到维度减小的新张量。

---

演示代码：

import torch

a=torch.arange(0,6)
b=a.view(2,3) # 维度：2*3
print("b",b)

c=b.unsqueeze(0) # 维度：1*2*3
print("c",c)

d=c.unsqueeze(2) # 维度：1*2*1*3
print("d",d)

e=d.squeeze(0) # 维度：2*1*3
print("e",e)

f=e.squeeze(0) # 维度：2*1*3
print("f",f)

结果显示：

b tensor([[0, 1, 2],
        [3, 4, 5]])
c tensor([[[0, 1, 2],
         [3, 4, 5]]])
d tensor([[[[0, 1, 2]],

         [[3, 4, 5]]]])
e tensor([[[0, 1, 2]],

        [[3, 4, 5]]])
f tensor([[[0, 1, 2]],

        [[3, 4, 5]]])