torch.squeeze() dim=1 dim=-1 dim=2

对数据的维度进行压缩

使用方式：torch.squeeze(input, dim=None, out=None)

将输入张量形状中的1 去除并返回。 如果输入是形如(A×1×B×1×C×1×D)，那么输出形状就为： (A×B×C×D)

import torch
x = torch.rand(2, 1, 1, 3, 1, 4)
print('=======x=========')
print(x.shape)
out_1 = torch.squeeze(x)
print('=======out_1=========')
print(out_1.shape)
# =======x=========
# torch.Size([2, 1, 1, 3, 1, 4])
# =======out_1=========
# torch.Size([2, 3, 4])

当给定dim时，那么挤压操作只在给定维度上。例如，输入形状为: (A×1×B), squeeze(input, 0) 将会保持张量不变，只有用 squeeze(input, 1)，形状会变成 (A×B)。

注意：

如果dim指定的维度的值为1

第一种情况

import torch
x = torch.rand(2,1,1,3,1,4)
print('=======x=========')
print(x.shape)
out_1 = torch.squeeze(x, dim=1)
print('=======out_1=========')
print(out_1.shape)
# =======x=========
# torch.Size([2, 1, 1, 3, 1, 4])
# =======out_1=========
# torch.Size([2, 1, 3, 1, 4])

import torch
x = torch.rand(2,1,3,1,4)
print('=======x=========')
print(x.shape)
out_1 = torch.squeeze(x, dim=1)
print('=======out_1=========')
print(out_1.shape)
# =======x=========
# torch.Size([2, 1, 3, 1, 4])
# =======out_1=========
# torch.Size([2, 3, 1, 4])

第二种情况

x = torch.rand(1,2,1,1,3,1,4)
print('=======x=========')
print(x.shape)
out_2 = torch.squeeze(x, dim=1)
print('=======out_2=========')
print(out_2.shape)
# =======x=========
# torch.Size([1, 2, 1, 1, 3, 1, 4])
# =======out_2=========
# torch.Size([1, 2, 1, 1, 3, 1, 4])

第三种情况

x = torch.rand(1,1,2,1,1,3,1,4)
print('=======x=========')
print(x.shape)
out_3 = torch.squeeze(x, dim=1)
print('=======out_3=========')
print(out_3.shape)
# =======x=========
# # torch.Size([1, 1, 2, 1, 1, 3, 1, 4])
# # =======out_3=========
# # torch.Size([1, 2, 1, 1, 3, 1, 4])

如果dim指定的维度的值为-1

第一种情况 如果dim指定的维度的值为-1

import torch

x = torch.rand(2,1,1,3,1,4)
print('=======x=========')
print(x.shape)
out_1 = torch.squeeze(x, dim=-1)
print('=======out_1=========')
print(out_1.shape)
# =======x=========
# torch.Size([2, 1, 1, 3, 1, 4])
# =======out_1=========
# torch.Size([2, 1, 1, 3, 1, 4])

第二种情况 如果dim指定的维度的值为-1

x = torch.rand(2,1,1,3,1,4,1)
print('=======x=========')
print(x.shape)
out_2 = torch.squeeze(x, dim=-1)
print('=======out_2=========')
print(out_2.shape)
# =======x=========
# torch.Size([2, 1, 1, 3, 1, 4, 1])
# =======out_2=========
# torch.Size([2, 1, 1, 3, 1, 4])

第三种情况 如果dim指定的维度的值为-1

x = torch.rand(2,1,1,3,1,4,1,1)
print('=======x=========')
print(x.shape)
out_3 = torch.squeeze(x, dim=-1)
print('=======out_3=========')
print(out_3.shape)
# =======x=========
# torch.Size([2, 1, 1, 3, 1, 4, 1, 1])
# =======out_3=========
# torch.Size([2, 1, 1, 3, 1, 4, 1])

如果dim指定的维度的值为2

import torch

x = torch.rand(2,1,3,1,4)
print('=======x=========')
print(x.shape)
out_1 = torch.squeeze(x, dim=2)
print('=======out_1=========')
print(out_1.shape)
# =======x=========
# torch.Size([2, 1, 3, 1, 4])
# =======out_1=========
# torch.Size([2, 1, 3, 1, 4])

x = torch.rand(2,1,1,3,1,4)
print('=======x=========')
print(x.shape)
out_2 = torch.squeeze(x, dim=2)
print('=======out_2=========')
print(out_2.shape)
# =======x=========
# torch.Size([2, 1, 1, 3, 1, 4])
# =======out_2=========
# torch.Size([2, 1, 3, 1, 4])


x = torch.rand(1,2,1,1,3,1,1,4)
print('=======x=========')
print(x.shape)
out_3 = torch.squeeze(x, dim=2)
print('=======out_3=========')
print(out_3.shape)
# =======x=========
# torch.Size([1, 2, 1, 1, 3, 1, 1, 4])
# =======out_3=========
# torch.Size([1, 2, 1, 3, 1, 1, 4])