Pytorch基础：Tensor的squeeze和unsqueeze方法

日晨难再

已于 2025-04-07 01:57:06 修改

阅读量1.9k

点赞数 19

分类专栏： Pytorch基础文章标签： pytorch 人工智能 python

于 2024-07-26 21:32:43 首次发布

本文链接：https://blog.csdn.net/weixin_45791458/article/details/140699760

版权

Pytorch基础专栏收录该内容

15 篇文章

订阅专栏

在Pytorch中，squeeze和unsqueeze是Tensor类的一个重要方法，同时它们也是torch模块中的一个函数，它们的语法如下所示。

Tensor.squeeze(dim=None) → Tensor
torch.squeeze(input, dim=None) → Tensor

input (Tensor) – the input tensor.
dim (int or tuple of ints, optional) – if given, the input will be squeezed only in the specified dimensions.


Tensor.unsqueeze(dim) → Tensor
torch.unsqueeze(input, dim) → Tensor

input (Tensor) – the input tensor.
dim (int) – the index at which to insert the singleton dimension

一、squeeze

squeeze函数（或方法）返回一个新的张量，该张量移除了原张量中大小为1的维度，例如：输入张量的形状是(A×1×B×C×1×D)，使用了squeeze函数（或方法）后，输出张量的形状是(A×B×C×D)。请注意：输出张量将与输入张量共享底层存储，因此改变一个张量的内容将改变另一个张量的内容。默认情况下，squeeze将移除所有尺寸为1的维度，如果传递了dim参数，则会将dim中的维度展开。dim的范围可以是[-input.dim()-1, input.dim()]，其中负数索引表示从后往前数的位置，例如-1代表最后一个维度。

可以看下面的例子以更好的理解：

import torch

# 创建一个形状为 (2, 1, 2, 1, 2) 的张量
x = torch.zeros(2, 1, 2, 1, 2)
print(x, x.size(), id(x))

# 移除所有大小为1的维度
a = torch.squeeze(x)  # 等价于 a = x.squeeze()
print(a, a.size(), id(a))

# 尝试移除第0维度（由于第0维度大小不为1，因此不改变形状）
b = torch.squeeze(x, 0)  # 等价于 b = x.squeeze(0)
print(b, b.size(), id(b))

# 移除第1维度（第1维度大小为1）
c = torch.squeeze(x, 1)  # 等价于 c = x.squeeze(1)
print(c, c.size(), id(c))

# 移除第1、第2和第3维度（第1和第3维度大小为1，第2维度不变）
d = torch.squeeze(x, (1, 2, 3))  # 等价于 d = x.squeeze((1, 2, 3))
print(d, d.size(), id(d))

# 验证所有张量共享底层存储空间
print(x.storage().data_ptr() == a.storage().data_ptr() == b.storage().data_ptr() == c.storage().data_ptr() == d.storage().data_ptr()) # 共享底层存储空间


输出：
tensor([[[[[0., 0.]],
          [[0., 0.]]]],
        [[[[0., 0.]],
          [[0., 0.]]]]]) torch.Size([2, 1, 2, 1, 2]) 1899057117680

tensor([[[0., 0.],
         [0., 0.]],
        [[0., 0.],
         [0., 0.]]]) torch.Size([2, 2, 2]) 1899057158240

tensor([[[[[0., 0.]],
          [[0., 0.]]]],
        [[[[0., 0.]],
          [[0., 0.]]]]]) torch.Size([2, 1, 2, 1, 2]) 1899737467296

tensor([[[[0., 0.]],
         [[0., 0.]]],
        [[[0., 0.]],
         [[0., 0.]]]]) torch.Size([2, 2, 1, 2]) 1899737467376

tensor([[[0., 0.],
         [0., 0.]],

        [[0., 0.],
         [0., 0.]]]) torch.Size([2, 2, 2]) 1899737467216
True

二、 unsqueeze

unsqueeze函数（或方法）函数返回一个新的张量，该张量在指定维度(dim)插入一个大小为1的维度。使用unsqueeze函数（或方法）后，输入张量的形状会相应增加一个维度。例如，输入张量的形状是（A×B×C），在第1维度使用unsqueeze后，输出张量的形状将变为（A×1×B×C）。请注意，输出张量将与输入张量共享底层存储，因此改变一个张量的内容将改变另一个张量的内容。dim的范围可以是[-input.dim(), input.dim()-1]，其中负数索引表示从后往前数的位置，例如-1代表最后一个维度。

可以看下面的例子以更好的理解：

import torch

# 创建一个形状为 (2, 2, 2) 的张量
x = torch.zeros(2, 2, 2)
print(x, x.size(), id(x))

# 在第0维度插入单维度
a = torch.unsqueeze(x, 0)  # 等价于 a = x.unsqueeze(0)
print(a, a.size(), id(a))

# 在第1维度插入单维度
b = torch.unsqueeze(x, 1)  # 等价于 b = x.unsqueeze(1)
print(b, b.size(), id(b))

# 在第2维度插入单维度
c = torch.unsqueeze(x, 2)  # 等价于 c = x.unsqueeze(2)
print(c, c.size(), id(c))

# 在第3维度插入单维度
d = torch.unsqueeze(x, 3)  # 等价于 d = x.unsqueeze(3)
print(d, d.size(), id(d))

# 验证所有张量共享底层存储空间
print(x.storage().data_ptr() == a.storage().data_ptr() == b.storage().data_ptr() == c.storage().data_ptr() == d.storage().data_ptr())  # 共享底层存储空间


输出：
tensor([[[0., 0.],
         [0., 0.]],
        [[0., 0.],
         [0., 0.]]]) torch.Size([2, 2, 2]) 1509028592032

tensor([[[[0., 0.],
          [0., 0.]],
         [[0., 0.],
          [0., 0.]]]]) torch.Size([1, 2, 2, 2]) 1509028632592

tensor([[[[0., 0.],
          [0., 0.]]],
        [[[0., 0.],
          [0., 0.]]]]) torch.Size([2, 1, 2, 2]) 1507561225888

tensor([[[[0., 0.]],
         [[0., 0.]]],
        [[[0., 0.]],
         [[0., 0.]]]]) torch.Size([2, 2, 1, 2]) 1507561391824

tensor([[[[0.],
          [0.]],
         [[0.],
          [0.]]],
        [[[0.],
          [0.]],
         [[0.],
          [0.]]]]) torch.Size([2, 2, 2, 1]) 1507561391904
True

三、squeeze_和unsqueeze_

Tensor还有拥有squeeze_和unsqueeze_方法，它们是原地操作(in-place)版本的，意思是它们会直接对原始张量进行操作，而不是创建新的张量。squeeze_和unsqueeze_的语法和功能与squeeze和unsqueeze类似，只是它们会直接修改调用它们的张量本身，除此之外它返回的也是修改后的变量本身（因此无需通过返回值传递常量）。由于是原地操作，squeeze_和unsqueeze_会节省内存，但是要小心使用，因为它们会改变原始数据。以下是一些例子，帮助理解这两个原地操作方法的使用：

import torch

# 创建一个形状为 (2, 1, 2, 1, 2) 的张量
x = torch.zeros(2, 1, 2, 1, 2)
print(x, x.size(), id(x))

# 原地移除所有大小为1的维度
x.squeeze_()  # 等价于 x = x.squeeze_()
print(x, x.size(), id(x))

# 重新创建一个形状为 (2, 1, 2, 1, 2) 的张量
x = torch.zeros(2, 1, 2, 1, 2)
print(x, x.size(), id(x))

# 原地移除第1维度（第1维度大小为1）
x.squeeze_(1)  # 等价于 x = x.squeeze_(1)
print(x, x.size(), id(x))

# 重新创建一个形状为 (2, 2, 2) 的张量
x = torch.zeros(2, 2, 2)
print(x, x.size(), id(x))

# 原地在第1维度插入单维度
x.unsqueeze_(1)  # 等价于 x = x.unsqueeze_(1)
print(x, x.size(), id(x))

输出：
tensor([[[[[0., 0.]],
          [[0., 0.]]]],
        [[[[0., 0.]],
          [[0., 0.]]]]]) torch.Size([2, 1, 2, 1, 2]) 140209009511456

tensor([[[0., 0.],
         [0., 0.]],
        [[0., 0.],
         [0., 0.]]]) torch.Size([2, 2, 2]) 140209009511456

tensor([[[[[0., 0.]],
          [[0., 0.]]]],
        [[[[0., 0.]],
          [[0., 0.]]]]]) torch.Size([2, 1, 2, 1, 2]) 140204776090096

tensor([[[[0., 0.]],
         [[0., 0.]]],
        [[[0., 0.]],
         [[0., 0.]]]]) torch.Size([2, 2, 1, 2]) 140204776090096

tensor([[[0., 0.],
         [0., 0.]],
        [[0., 0.],
         [0., 0.]]]) torch.Size([2, 2, 2]) 140209009511455

tensor([[[[0., 0.],
          [0., 0.]]],
        [[[0., 0.],
          [0., 0.]]]]) torch.Size([2, 1, 2, 2]) 140209009511455