torch.squeeze与torch.unsqueeze——数组的压缩与解压

torch.squeeze()

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

功能：将尺寸大小为1的维度进行删除操作。

输入：

• input：需要删除的张量数组
• dim：在指定维度上进行删除操作

注意：

• 若不指定dim，则默认将所有尺寸大小为1的维度删除；若指定dim，则只在指定的维度上进行删除操作。
• 返回的张量与输入的张量共享存储，因此更改一个数组的数据内容，将会更改另一个数组。
• torch.squeeze与a.squeeze使用效果一样，只是后者默认在数组a上进行操作。

代码案例

import torch
a=torch.arange(5).view(1,5,1)
b=torch.squeeze(a)
print(a)
print(b)
print(a.shape)
print(b.shape)

输出

# 压缩前
tensor([[[0],
         [1],
         [2],
         [3],
         [4]]])
# 压缩后
tensor([0, 1, 2, 3, 4])
# 压缩前尺寸
torch.Size([1, 5, 1])
# 压缩后尺寸
torch.Size([5])

可见，若不指定dim，则会将所有尺寸大小为1的维度全部删除。

指定dim时，只会在该dim上进行操作。

import torch
a=torch.arange(5).view(5,1)
b=a.squeeze(dim=0)
c=a.squeeze(dim=1)
print(a)
print(b)
print(c)
print(a.shape)
print(b.shape)
print(c.shape)
# 存储共享测试
c[0]=9
print(a)
print(c)

输出

# 原数组
tensor([[0],
        [1],
        [2],
        [3],
        [4]])
# dim设为0时
tensor([[0],
        [1],
        [2],
        [3],
        [4]])
# dim设为1时
tensor([0, 1, 2, 3, 4])
# 原数组尺寸
torch.Size([5, 1])
# dim设为0时尺寸
torch.Size([5, 1])
# dim设为1时尺寸
torch.Size([5])

# 当c被修改时，a相应位置的数据也会被修改
tensor([[9],
        [1],
        [2],
        [3],
        [4]])
tensor([9, 1, 2, 3, 4])

如果指定的维度上尺寸不是1，则不进行操作；如果指定的维度上尺寸是1，则只删除该维度。压缩后的数据如果修改了某个数据，则也会在原数组相应位置进行修改。

torch.unsqueeze()

torch.unsqueeze(input, dim) → Tensor

功能：在指定的维度上，插入尺寸大小为1的维度(相当于是squeeze的逆运算)

输入：

• input：输入张量
• dim：指定的维度

注意：

• 与squeeze类似，返回的张量与输入的张量共享存储，因此更改一个数组的数据内容，将会更改另一个数组。
• torch.unsqueeze与a.unsqueeze使用效果一样，只是后者默认在数组a上进行操作。
• 必须要指定dim，这里与squeeze不同。

代码案例

import torch
a=torch.arange(5).view(5)
b=torch.unsqueeze(a,dim=0)
c=torch.unsqueeze(a,dim=1)
print(a)
print(b)
print(c)
print(a.shape)
print(b.shape)
print(c.shape)

输出

# 原数组
tensor([0, 1, 2, 3, 4])
# dim设为0时
tensor([[0, 1, 2, 3, 4]])
# dim设为1时
tensor([[0],
        [1],
        [2],
        [3],
        [4]])
# 原数组尺寸
torch.Size([5])
# dim设为0时尺寸
torch.Size([1, 5])
# dim设为1时尺寸
torch.Size([5, 1])

官方文档

torch.squeeze()：https://pytorch.org/docs/stable/generated/torch.squeeze.html?highlight=torch%20squeeze#torch.squeeze

torch.unsqueeze()：https://pytorch.org/docs/stable/generated/torch.unsqueeze.html?highlight=torch%20unsqueeze#torch.unsqueeze

点个赞支持一下吧