总而言之,squeeze是用来压缩一些轴(维度为1)的,相反,unsqueeze用来增加张量轴的数量的。
squeeze()
import torch
a = torch.arange(10).reshape((2,1,5,1))
a
-----out
tensor([[[[0],
[1],
[2],
[3],
[4]]],
[[[5],
[6],
[7],
[8],
[9]]]])
b = a.squeeze()
b
-----out
tensor([[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9]])
以上可以看出,张量从4个轴,变成了2个轴,其中维度为1的被压缩了。
unsqueeze()
c = b.unsqueeze(1)
c
-----out
tensor([[[0, 1, 2, 3, 4]],
[[5, 6, 7, 8, 9]]])
c = b.unsqueeze(2)
c
-----out
tensor([[[0],
[1],
[2],
[3],
[4]],
[[5],
[6],
[7],
[8],
[9]]])
其中数字代表需要增加的轴。(其实具体的轴我也转不清,在用到的时候应该就知道了。)
目前比较会用到的场景是
d = torch.arange(6).reshape((1,-1))
d
-----out
tensor([[0, 1, 2, 3, 4, 5]])
d.squeeze()
-----out
tensor([0, 1, 2, 3, 4, 5])
tensor([[0, 1, 2, 3, 4, 5]])
和 tensor([0, 1, 2, 3, 4, 5])
看上去是一样的,但是在pytorch里,前者是矩阵,后者是向量,在计算的过程中,要是搞不清楚,还是比较麻烦的。之前我会比较笨,用reshape的方式,现在看来,其实只要用squeeze的方法就可以了。