Adding a dimension to a tensor can be important when you’re building machine learning models. Although the actual PyTorch function is called
unsqueeze(), you can think of this as the PyTorch “add dimension” operation. Let’s look at two ways to do it.
The easiest way to expand tensors with dummy dimensions is by inserting
None into the axis you want to add. For example, say you have a feature vector with 16 elements. To add a dummy batch dimension, you should index the 0th axis with
import torch x = torch.randn(16) x = x[None, :] x.shape # Expected result # torch.Size([1, 16])
The slicing syntax works by specifying new dimensions with
None and existing dimensions with a colon. That means you can prepend more than one dimension if you want:
x = torch.randn(16) x = x[None, None, :] x.shape # Expected result # torch.Size([1, 1, 16])
Additionally, when you prepend indices, you don’t actually have to specify the colons. All the remaining dimensions are implied. That means you can repeat the same example with no colon at the end:
x = torch.randn(16) x = x[None, None] x.shape # Expected result # torch.Size([1, 1, 16])
This is a very simple trick for prepending axes to the front of tensors.
You can also add dimensions to the end of tensors, which can be useful for broadcasting operations like pairwise distance. All you have to do is rearrange the colons and the
Let’s convert a 2D tensor with shape
(3, 4) to a 3D tensor with shape
(3, 4, 1):
x = torch.randn(3, 4) x = x[:, :, None] x.shape # Expected result # torch.Size([3, 4, 1])
Notice that we have to specify the first two dimensions with colons and then add a
None to the end. When you append dimensions to the end, the colons are required.
And naturally, this trick works regardless of where you want to insert the dimension. For example, if you wanted to add a dimension to the middle of the previous tensor, you could write
x[:, None, :].
Overall, I prefer using
None to add dimensions because the same syntax works in NumPy:
import numpy as np x = np.random.randn(16) x = x[None] x.shape # Expected result # (1, 16)
But there is another approach that you should know about. PyTorch provides a function called
unsqueeze() that does the same thing.
x = torch.randn(16) x = torch.unsqueeze(x, dim=0) x.shape # Expected result # torch.Size([1, 16])
dim argument is how you specify where the new axis should go. To put a new dimension on the end, pass
x = torch.randn(3, 4) x = torch.unsqueeze(x, dim=-1) x.shape # Expected result # torch.Size([3, 4, 1])
Not bad. But you have to be careful if you use both NumPy and PyTorch because there is no NumPy
x = np.random.randn(16) try: x = np.unsqueeze(x, axis=0) except AttributeError as e: print(e) # Expected error message # module 'numpy' has no attribute 'unsqueeze'
Indeed, the equivalent NumPy function is called
x = np.random.randn(16) x = np.expand_dims(x, axis=0) x.shape # Expected error # (1, 16)
Better to use
None indexing for the PyTorch “add dimension” op to avoid confusion.