SoFunction
Updated on 2025-04-10

Extra detailed explanation of squeeze in python (with code example)

Squeeze operation in Python

Squeeze is a forRemoves dimensions of size 1 in tensor or arrayOperation.

It can be found inPyTorchandNumPyUsed in  . In practical applications,squeezeOperations are often used to adjust the shape of data to meet the needs of a specific operation or model.

Main functions:

  • Removes an axis with dimension 1: For example, if a tensor has a shape of(1, 3, 1), usesqueezeIt will become(3,), i.e. all dimensions of size 1 are removed.
  • Stay non-1 dimension：squeezeOnly dimensions with size 1 are removed, and other dimensions will not change.

squeeze in PyTorch

In PyTorch,squeeze()Used to remove all or specified single dimensions in a tensor (dimensions of size 1).

The syntax is as follows:

(input, dim=None)
  • input: The input tensor.
  • dim(Optional): Specifies the dimension to be removed. If the dimension is specified and the dimension has a size of 1, then removes the dimension; if not specified, all dimensions with a size of 1 are removed by default.

Example 1: Remove all single dimensions

import torch

# Create a tensor with shape (1, 3, 1)x = ([[[1], [2], [3]]])
print("Original shape:", )

# Use squeeze to remove all dimensions with 1x_squeezed = (x)
print("Squeezed shape:", x_squeezed.shape)

Output：

Original shape: ([1, 3, 1])
Squeezed shape: ([3])

explain：

  • The shape of the original tensor is(1, 3, 1),Right nowThe first dimensionandThe last dimensionThe size is 1.
  • squeeze()After that, all dimensions of size 1 are removed and the resulting tensor shape becomes(3),Right nowRemoved the first dimensionandThe last dimension。

Example 2: Specify the removal dimension

# Create a tensor with shape (1, 3, 1)x = ([[[1], [2], [3]]])

# Use squeeze to remove dimension 0 (if the dimension size is 1)x_squeezed = (x, dim=0)
print("Squeezed shape:", x_squeezed.shape)

Output：

Squeezed shape: ([3, 1])

explain：

  • This is specifieddim=0, indicating the removal of dimension 0 (size 1). In this way, the shape of the tensor is from(1, 3, 1)Become(3, 1)。
  • If you specifieddim=2, but the size of this dimension is not 1, so the dimension will not be removed.

squeeze in NumPy

existNumPymiddle,squeeze()There is a similar function to remove all dimensions in the array with a specified size of 1. The syntax is as follows:

(a, axis=None)
  • a: The input array.
  • axis(Optional): Specifies the dimension to be removed, if the specified dimension size is 1, remove the dimension; if not specified, remove all dimensions with a size of 1.

Example 1: Remove all single dimensions

import numpy as np

# Create an array of shapes (1, 3, 1)x = ([[[1], [2], [3]]])
print("Original shape:", )

# Use squeeze to remove all dimensions with 1x_squeezed = (x)
print("Squeezed shape:", x_squeezed.shape)

Output：

Original shape: (1, 3, 1)
Squeezed shape: (3,)

explain：

  • The shape of the original array is(1, 3, 1), where the size of the first and third dimensions is 1.
  • usesqueeze()After that, all dimensions of size 1 are removed, and the shape is finally obtained(3,)array.

Example 2: Specify the removal dimension

# Create an array of shapes (1, 3, 1)x = ([[[1], [2], [3]]])

# Use squeeze to remove dimension 0x_squeezed = (x, axis=0)
print("Squeezed shape:", x_squeezed.shape)

Output：

Squeezed shape: (3, 1)

explain：

  • Designatedaxis=0, indicating the removal of dimension 0 (size 1). Therefore, the shape of the tensor is from(1, 3, 1)Become(3, 1)。

When to use squeeze?

  • Remove redundant dimensions: When a tensor or array contains redundant dimensions (dimensionality of size 1), usesqueeze()Can simplify the data structure.
  • Adapt model input: In deep learning models, specific input dimensions are often required. If the dimensions of the data do not meet the requirements, you can usesqueeze()Remove unnecessary single dimensions.
  • Avoid inconsistency in dimensions: In some operations, some operations may produce unnecessary single dimensions, usingsqueeze()It can maintain dimensional consistency of data.

Summarize

  • squeezeUsed forRemoves dimensions of size 1 in tensor or array, simplify the data structure.
  • existPyTorchandNumPymiddle,squeeze()There are similar functions to remove all or specified dimensions of size 1.
  • squeeze()It is a common operation when processing data dimensions, adapting to model input or data storage.

By removing useless single dimensions, we can simplify the data shape to make it more suitable for subsequent processing and calculations.

This is the end of this article about the super detailed explanation of squeeze in python. For more related content of python squeeze, please search for my previous articles or continue browsing the related articles below. I hope everyone will support me in the future!