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.
Using None indexing
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 None:
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 None(s).
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)PyTorch Unsqueeze
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])The dim argument is how you specify where the new axis should go. To put a new dimension on the end, pass dim=-1:
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 unsqueeze() function:
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 expand_dims():
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.