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

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

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

函数描述:

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])

CJ:只有维度为1的才能被去掉

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值