Pytorch squeeze and unsqueeze

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?

3 Answers

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.

Simply put, unsqueeze() "adds" a superficial 1 dimension to tensor (at the specified dimension), while squeeze removes all superficial 1 dimensions from tensor.

You should look at tensor's shape attribute to see it easily. In your last case it would be:

import torch tensor = torch.tensor([1, 0, 2, 3, 4]) tensor.shape # torch.Size([5]) tensor.unsqueeze(dim=0).shape # [1, 5] tensor.unsqueeze(dim=1).shape # [5, 1] 

It is useful for providing single sample to the network (which requires first dimension to be batch), for images it would be:

# 3 channels, 32 width, 32 height tensor = torch.randn(3, 32, 32) # 1 batch, 3 channels, 32 width, 32 height tensor.unsqueeze(dim=0).shape 

unsqueeze can be seen if you create tensor with 1 dimensions, e.g. like this:

# 3 channels, 32 width, 32 height and some 1 unnecessary dimensions tensor = torch.randn(3, 1, 32, 1, 32, 1) # 1 batch, 3 channels, 32 width, 32 height again tensor.squeeze().unsqueeze(0) # [1, 3, 32, 32] 

1. torch.unsqueeze(input, dim) → Tensor

   a = torch.randn(4, 4, 4) torch.unsqueeze(a, 0).size() >>> torch.Size([1, 4, 4, 4]) 

   a = torch.randn(4, 4, 4) torch.unsqueeze(a, 1).size() >>> torch.Size([4, 1, 4, 4]) 

   a = torch.randn(4, 4, 4) torch.unsqueeze(a, 2).size() >>> torch.Size([4, 4, 1, 4]) 

   a = torch.randn(4, 4, 4) torch.unsqueeze(a, 3).size() >>> torch.Size([4, 4, 4, 1]) 

2. torch.squeeze(input, dim=None, out=None) → Tensor

   b = torch.randn(4, 1, 4) >>> tensor([[[ 1.2912, -1.9050, 1.4771, 1.5517]], [[-0.3359, -0.2381, -0.3590, 0.0406]], [[-0.2460, -0.2326, 0.4511, 0.7255]], [[-0.1456, -0.0857, -0.8443, 1.1423]]]) 

   b.size() >>> torch.Size([4, 1, 4]) 

   c = b.squeeze(1) 

   b >>> tensor([[[ 1.2912, -1.9050, 1.4771, 1.5517]], [[-0.3359, -0.2381, -0.3590, 0.0406]], [[-0.2460, -0.2326, 0.4511, 0.7255]], [[-0.1456, -0.0857, -0.8443, 1.1423]]]) 

   b.size() >>> torch.Size([4, 1, 4]) 

   c >>> tensor([[ 1.2912, -1.9050, 1.4771, 1.5517], [-0.3359, -0.2381, -0.3590, 0.0406], [-0.2460, -0.2326, 0.4511, 0.7255], [-0.1456, -0.0857, -0.8443, 1.1423]]) 

   c.size() >>> torch.Size([4, 4])