python中squeeze的超详细解释(附代码示例)-QQ沐编程

这篇文章主要介绍了python中squeeze的超详细解释,squeeze操作用于去除张量或数组中大小为1的维度,简化数据结构,在PyTorch和NumPy中都有类似的功能,需要的朋友可以参考下

Python 中的 squeeze 操作

Squeeze 是一个用于 去除张量或数组中大小为 1 的维度 的操作。

它可以在 PyTorch 和 NumPy 中使用。在实际应用中，squeeze 操作常用于调整数据的形状，以满足特定操作或模型的需求。

主要作用：

去除维度为 1 的轴：例如，如果一个张量的形状为 (1, 3, 1), 使用 squeeze 后会变成 (3,)，即去除了所有大小为 1 的维度。
保持非 1 维度：squeeze 只去除大小为 1 的维度，而其他维度不会改变。

PyTorch 中的 squeeze

在 PyTorch 中，squeeze() 用于去除张量中所有或指定的单维度（大小为 1 的维度）。

其语法如下：

torch.squeeze(input, dim=None)

input：输入的张量。
dim（可选）：指定要去除的维度，如果指定该维度并且该维度的大小为 1，则去除该维度；如果不指定，默认去除所有维度大小为 1 的维度。

示例 1：去除所有单维度

import torch

# 创建一个形状为 (1, 3, 1) 的张量
x = torch.tensor([[[1], [2], [3]]])
print("Original shape:", x.shape)

# 使用 squeeze 去除所有维度为 1 的维度
x_squeezed = torch.squeeze(x)
print("Squeezed shape:", x_squeezed.shape)

输出：

Original shape: torch.Size([1, 3, 1])
Squeezed shape: torch.Size([3])

解释：

原始张量的形状是 (1, 3, 1)，即第一个维度和最后一个维度的大小为 1。
squeeze() 后，所有大小为 1 的维度被去除，结果的张量形状变为 (3)，即去除了第一个维度和最后一个维度。

示例 2：指定去除维度

# 创建一个形状为 (1, 3, 1) 的张量
x = torch.tensor([[[1], [2], [3]]])

# 使用 squeeze 去除第 0 维（如果该维度大小为 1）
x_squeezed = torch.squeeze(x, dim=0)
print("Squeezed shape:", x_squeezed.shape)

输出：

Squeezed shape: torch.Size([3, 1])

解释：

这里指定了 dim=0，表示去除第 0 维（大小为 1）。这样，张量的形状从 (1, 3, 1) 变成了 (3, 1)。
如果你指定了 dim=2，但是该维度的大小不是 1，那么就不会去除该维度。

NumPy 中的 squeeze

在 NumPy 中，squeeze() 也有类似的功能，用于去除数组中所有或指定的大小为 1 的维度。其语法如下：

numpy.squeeze(a, axis=None)

a：输入的数组。
axis（可选）：指定要去除的维度，如果指定的维度大小为 1，则去除该维度；如果不指定，则去除所有大小为 1 的维度。

示例 1：去除所有单维度

import numpy as np

# 创建一个形状为 (1, 3, 1) 的数组
x = np.array([[[1], [2], [3]]])
print("Original shape:", x.shape)

# 使用 squeeze 去除所有维度为 1 的维度
x_squeezed = np.squeeze(x)
print("Squeezed shape:", x_squeezed.shape)

输出：

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

解释：

原始数组的形状是 (1, 3, 1)，其中第一个和第三个维度的大小为 1。
使用 squeeze() 后，所有大小为 1 的维度被去除，最终得到形状为 (3,) 的数组。

示例 2：指定去除维度

# 创建一个形状为 (1, 3, 1) 的数组
x = np.array([[[1], [2], [3]]])

# 使用 squeeze 去除第 0 维
x_squeezed = np.squeeze(x, axis=0)
print("Squeezed shape:", x_squeezed.shape)

输出：