Add A New Dimension To The End Of A Tensor In PyTorch

This video will show you how to add a new dimension to the end of a PyTorch tensor by using None-style indexing.

First, we import PyTorch.

import torch

Then we print the PyTorch version we are using.

print(torch.__version__)

We are using PyTorch 0.4.0.

Let’s now create a PyTorch tensor of size 2x4x6x8 using the PyTorch Tensor operation, and we want the dimensions to be 2x4x6x8.

pt_empty_tensor_ex = torch.Tensor(2,4,6,8)

This is going to return to us an uninitialized tensor which we assign to the Python variable pt_empty_tensor_ex.

Let’s check what dimensions our pt_empty_tensor_ex Python variable has.

print(pt_empty_tensor_ex.size())

We see that it is a 2x4x6x8 tensor.

What we want to do now is we want to add a new axis to the end of this tensor.

So it’ll be 2x4x6x8x1.

The way we’re going to do this is we’re going to use the None-style indexing.

So we pass in our initial tensor, pt_empty_tensor_ex, and then we’re going to do indexing to specify what it is that we want.

pt_extend_end_tensor_ex = pt_empty_tensor_ex[:,:,:,:,None]

So for the first index, we pass in a colon to specify that we want everything in the already existing first dimension.

For the second index, we use a colon to specify that we want everything in the already existing second dimension.

For the third index, we use a colon to specify that we want everything in the already existing third dimension.

For the fourth index, we specify with a colon that we want everything in the already existing fourth dimension.

Then finally, because we had four dimensions and now we’re going to add something new, we are going to use the None-style indexing with a capital N.

What we are specifying here is that we want to create a new axis right at the end of the tensor.

This tensor is then going to be assigned to the Python variable pt_extend_end_tensor_ex.

Let’s check the dimensions by using the PyTorch size operation to see if PyTorch did create a new axis for us at the end of the tensor.

print(pt_extend_end_tensor_ex.size())

Bingo! We see that we now have a new axis at the end.

Before, it was 2x4x6x8.

Now, it’s 2x4x6x8x1.

Perfect - We were able to add a new dimension to the end of a PyTorch tensor by using None style indexing.