Pytorch squeeze and unsqueeze

Question

I don't understand what squeeze and unsqueeze do to a tensor, even after looking at the docs and related questions.

I tried to understand it by exploring it myself in python. I first created a random tensor with

x = torch.rand(3,2,dtype=torch.float)
>>> x
tensor([[0.3703, 0.9588],
        [0.8064, 0.9716],
        [0.9585, 0.7860]])

But regardless of how I squeeze it, I end up with the same results:

torch.equal(x.squeeze(0), x.squeeze(1))
>>> True

If I now try to unsqueeze I get the following,

>>> x.unsqueeze(1)
tensor([[[0.3703, 0.9588]],
        [[0.8064, 0.9716]],
        [[0.9585, 0.7860]]])
>>> x.unsqueeze(0)
tensor([[[0.3703, 0.9588],
         [0.8064, 0.9716],
         [0.9585, 0.7860]]])
>>> x.unsqueeze(-1)
tensor([[[0.3703],
         [0.9588]],
        [[0.8064],
         [0.9716]],
        [[0.9585],
         [0.7860]]])

However if I now create a tensor x = torch.tensor([1,2,3,4]), and I try to unsqueeze it then it appears that 1 and -1 makes it a column where as 0 remains the same.

x.unsqueeze(0)
tensor([[1, 2, 3, 4]])
>>> x.unsqueeze(1)
tensor([[1],
        [2],
        [3],
        [4]])
>>> x.unsqueeze(-1)
tensor([[1],
        [2],
        [3],
        [4]])

Can someone provide an explanation of what squeeze and unsqueeze are doing to a tensor? And what's the difference between providing the arguements 0, 1 and -1?

Answer 1

Here is a visual representation of what squeeze/unsqueeze do for an effectively 2d matrix:

When you are unsqueezing a tensor, it is ambiguous which dimension you wish to 'unsqueeze' it across (as a row or column etc). The dim argument dictates this - i.e. position of the new dimension to be added.

Hence the resulting unsqueezed tensors have the same information, but the indices used to access them are different.