torch.unsqueeze()与torch.squeeze()用法

原创已于 2022-09-06 21:36:31 修改 · 421 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #机器学习

于 2022-09-03 21:50:01 首次发布

深度学习&Pytorch 专栏收录该内容

10 篇文章

订阅专栏

本文介绍了PyTorch中的unsqueeze和squeeze函数，用于调整Tensor的维度。unsqueeze能在指定位置插入一个维度，而squeeze则删除尺寸为1的维度。示例中展示了如何在不同位置增加和移除维度，并解释了只有尺寸为1的维度才能被squeeze函数删除。

部署运行你感兴趣的模型镜像

函数描述：

unsqueeze(input, dim) → Tensor

作用：在指定位置插入一个维度，对数据维度进行扩充

input：输入的Tensor
dim：要插入的维度

a = torch.arange(6).reshape(2, 3)
print(a)
b = a.unsqueeze(1)#在第2维度加一维度
print(b)
print(b.shape)
>>>
tensor([[0, 1, 2],
        [3, 4, 5]])
tensor([[[0, 1, 2]],

        [[3, 4, 5]]])
torch.Size([2, 1, 3])

函数描述：

squeeze(input, dim) → Tensor

作用：对数据维度进行压缩

a = torch.arange(12).reshape(1, 2, 6)
print(a)
a1 = a.squeeze(0)#将第一个维度去掉 
print(a1)
print(a1.shape)
>>>
tensor([[[ 0,  1,  2,  3,  4,  5],
         [ 6,  7,  8,  9, 10, 11]]])
tensor([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11]])
torch.Size([2, 6])
-----------------------
a2 = a.squeeze(-1)#最后一个维度并没有被去掉，因为不为1
print(a2)
print(a2.shape)
>>>
tensor([[[ 0,  1,  2,  3,  4,  5],
         [ 6,  7,  8,  9, 10, 11]]])
torch.Size([1, 2, 6])