pytorch api文档：张量的 squeeze()方法-删除维度

作者：马育民 • • 阅读：10039

介绍

PyTorch 中张量的 .squeeze() 方法，它正好是 .unsqueeze() 的逆操作，作用是 删除张量中所有大小为 1 的维度（也可以指定只删除某个位置的大小为 1 的维度），能简化张量的维度结构。

语法

torch.Tensor.squeeze(dim=None)

参数解释

  • 若不指定 dim：删除张量 所有shape大小为 1 的维度；
  • 若指定 dim：仅删除该维度（前提是这个维度的大小为 1，否则不生效，也不报错）；

返回值

返回新张量，原张量不会被修改（squeeze_() 是原地修改版本）。

注意事项

  1. squeeze() 只删除 大小为 1 的维度，非 1 的维度不会被改变；
  2. 不指定 dim 时，会删除所有大小为 1 的维度，需谨慎（避免误删需要保留的 1 维）；
  3. 原地修改用 squeeze_()（带下划线），比如 x.squeeze_() 直接修改 x 本身；
  4. PyTorch 也支持 torch.squeeze(x, dim) 写法，和 x.squeeze(dim) 效果一致。

例子

不指定 dim

import torch

# 示例1：一维张量（无大小为1的维度，squeeze后不变）
x1 = torch.tensor([1, 2, 3])  # shape [3]
print("x1 原始shape：", x1.shape)  # torch.Size([3])
print("x1 squeeze后shape：", x1.squeeze().shape)  # torch.Size([3])

# 示例2：含多个大小为1的维度
x2 = torch.tensor([[[1, 2, 3]]])  # shape [1, 1, 3]
print("\nx2 原始shape：", x2.shape)  # torch.Size([1, 1, 3])
x2_squeezed = x2.squeeze()
print("x2 squeeze后shape：", x2_squeezed.shape)  # torch.Size([3])
print("x2 squeeze后值：", x2_squeezed)  # tensor([1, 2, 3])

# 示例3：部分维度为1
x3 = torch.randn(2, 1, 4, 1)  # shape [2, 1, 4, 1]
print("\nx3 原始shape：", x3.shape)  # torch.Size([2, 1, 4, 1])
print("x3 squeeze后shape：", x3.squeeze().shape)  # torch.Size([2, 4])

输出结果：

x1 原始shape： torch.Size([3])
x1 squeeze后shape： torch.Size([3])

x2 原始shape： torch.Size([1, 1, 3])
x2 squeeze后shape： torch.Size([3])
x2 squeeze后值： tensor([1, 2, 3])

x3 原始shape： torch.Size([2, 1, 4, 1])
x3 squeeze后shape： torch.Size([2, 4])

指定 dim 的用法

当你只想删除某一个特定的、大小为 1 的维度 时，传入 dim 参数：

# 张量 shape [1, 3, 1, 5]
x4 = torch.randn(1, 3, 1, 5)
print("x4 原始shape：", x4.shape)  # torch.Size([1, 3, 1, 5])

# 只删除 dim=0 的维度（大小为1）
x4_dim0 = x4.squeeze(dim=0)
print("删除 dim=0 后shape：", x4_dim0.shape)  # torch.Size([3, 1, 5])

# 只删除 dim=2 的维度（大小为1）
x4_dim2 = x4_dim0.squeeze(dim=2)
print("删除 dim=2 后shape：", x4_dim2.shape)  # torch.Size([3, 5])

# 尝试删除非1的维度（dim=1，大小为3），无效果
x4_dim1 = x4.squeeze(dim=1)
print("删除 dim=1 后shape：", x4_dim1.shape)  # torch.Size([1, 3, 1, 5])

负索引支持

dim 也可以用负索引（-1 表示最后一个维度），规则和 unsqueeze() 一致：

x5 = torch.randn(2, 1)  # shape [2, 1]
# 删除最后一个维度（dim=-1，大小为1）
x5_squeezed = x5.squeeze(dim=-1)
print("x5 squeeze(-1)后shape：", x5_squeezed.shape)  # torch.Size([2])

使用场景

去除批量维度（batch dim）

模型输出通常带 batch 维度（比如 [1, 10] 表示 1 个样本的 10 类概率），用 squeeze() 去掉多余的 batch 维度：

# 模型输出：shape [1, 10]（1个样本，10类）
output = torch.randn(1, 10)
# 去除batch维度，变成 [10]
pred = output.squeeze(0)
print("预测值维度：", pred.shape)  # torch.Size([10])

清理维度，方便后续计算

比如从 [2, 1, 4] 清理为 [2, 4]，避免维度冗余导致的运算麻烦：

a = torch.randn(2, 1, 4)
b = torch.randn(2, 4)
# 直接相加会报错（维度不匹配），先squeeze
c = a.squeeze(1) + b
print("相加结果shape：", c.shape)  # torch.Size([2, 4])

和 unsqueeze() 配合使用

先增维适配运算，再删维还原结构：

# 原始张量 [3]
x = torch.tensor([1, 2, 3])
# 增维：[3, 1]
x_unsq = x.unsqueeze(1)
# 运算：和 [1, 3] 广播相加
res = x_unsq + torch.tensor([[4,5,6]])  # shape [3, 3]
# 若需要还原，可针对特定维度squeeze（这里无1维，squeeze无效果）
print(res.squeeze().shape)  # torch.Size([3, 3])

总结

  1. .squeeze(dim) 核心作用是删除张量中大小为 1 的维度（指定 dim 则仅删该维度，前提是其大小为 1）；
  2. 常用场景：去除多余的 batch 维度、清理冗余维度、配合 unsqueeze() 完成维度适配；
  3. 注意仅删除大小为 1 的维度，不指定 dim 会删除所有此类维度，原地修改用 squeeze_()。

原文出处：/show_1GW2esF8xpDO.html

吉ICP备15003043号