Adding a dimension to a tensor can be important when you’re building
deep learning models. In NumPy, you can do this by inserting None
into
the axis you want to add:
import numpy as np
x1 = np.zeros((10, 10))
x2 = x1[None, :, :]
>>> print(x2.shape)
(1, 10, 10)
EDIT: Good news! As of version 0.1.10, PyTorch supports None
-style indexing. You should probably use that. But if you prefer to do it the old-fashioned way, read on.
Fortunately, it’s easy enough in PyTorch. Just pass the axis index into
the .unsqueeze()
method.
import torch
x1 = torch.zeros(10, 10)
x2 = x1.unsqueeze(0)
>>> print(x2.size())
torch.Size([1, 10, 10])
You can also do it in place using the underscore version
.unsqueeze_()
.
x1 = torch.zeros(10, 10)
x1.unsqueeze_(0)
>>> print(x1.size())
torch.Size([1, 10, 10])
Want to know more cool tensor tricks? Check out my post on pairwise distance!