PyTorch torch.unsqueeze() Method

by

PyTorch torch.unsqueeze() method “returns a new tensor with a dimension of size one inserted at the specified position.”

The returned tensor shares the same underlying data with this tensor.

A dim value within the range [-input.dim() – 1, input.dim() + 1) can be used. Negative dim will correspond to unsqueeze() applied at dim = dim + input.dim() + 1.

Syntax

torch.unsqueeze(input, dim)

Parameters

  1. input (Tensor): It is an input tensor.
  2. dim (int): It is the index to insert the singleton dimension.

Return value

This function returns a new tensor with a dimension of size one inserted at the specified position. The returned tensor shares the same underlying data as the original tensor.

Example 1: How to Use torch.unsqueeze() method

import torch

# create a tensor
x = torch.tensor([1, 2, 3, 4])

print("Original tensor:")
print(x)
print("Shape of original tensor:", x.shape)

y = torch.unsqueeze(x, 0)

print("\nTensor after unsqueeze:")
print(y)
print("Shape of tensor after unsqueeze:", y.shape)

Output

Original tensor:
tensor([1, 2, 3, 4])
Shape of original tensor: torch.Size([4])

Tensor after unsqueeze:
tensor([[1, 2, 3, 4]])
Shape of tensor after unsqueeze: torch.Size([1, 4])

In this example, torch.unsqueeze(x, 0) is used to add an extra dimension at the 0th position. If x were of shape (N,), y would be of shape (1, N) after unsqueezing.

Example 2: Adding an extra dimension at the 1st position of the 2D Tensor

import torch

# create a 2D tensor
x = torch.tensor([[1, 2, 3], [4, 5, 6]])

print("Original tensor:")
print(x)
print("Shape of original tensor:", x.shape)

# add an extra dimension at the 1st position
y = torch.unsqueeze(x, 1)

print("\nTensor after unsqueeze:")
print(y)
print("Shape of tensor after unsqueeze:", y.shape)

# add an extra dimension at the last position
z = torch.unsqueeze(x, -1)

print("\nTensor after another unsqueeze:")
print(z)
print("Shape of tensor after another unsqueeze:", z.shape)

Output

PyTorch torch.unsqueeze()

 

As you can see, torch.unsqueeze() can add an extra dimension at any position of the tensor. The size of this new dimension is always 1. This can be useful in many situations, for example, when you need to match the shapes of different tensors for an operation requiring them to have the same dimensions.

That’s it.

Related posts

torch.squeeze()

torch.clone()

torch.nn.Conv2d()

torch.sum()

torch.full()

Post Views: 2

Leave a Comment