PyTorch：torch.Tensor.unsqueeze()、squeeze()

目录

1、unsqueeze()

2、squeeze()

 

 

---

1、unsqueeze()

作用：给指定的tensor增加一个指定(之前不存在的)的维度。通常用在两tensor相加，但不满足shape一致，同时又不符合广播机制的情况下。

举例说明：

# 两tensor相加，但不满足shape一致，同时又不符合广播机制的例子：
a = torch.ones(3,2,4)  # 创建3*2*4的全为1的tensor
b = torch.zeros(3,4)   # 创建3*4的全为0的tensor
c = a + b   # 直接相加 这样相加会报错
print('a:',a)
print('b:',b)
print('c:',c)

'''   报错如下   '''
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1

'''==================================================================================='''

# 两tensor相加，但不满足shape一致，使用unsqueeze对b进行维度扩充，使其满足广播机制的例子：
a = torch.ones(3,2,4)  # 创建3*2*4的全为1的tensor
b = torch.zeros(3,4)   # 创建3*4的全为0的tensor
b = b.unsqueeze(1) # 在增添一个维度1，即将b的shape转换为（3，1，4）
c = a + b   # 相加 这次就可以相加了
print('a:', a)
print('b:', b)
print('c:', c)

'''   结果如下   '''
a: tensor([[[1., 1., 1., 1.],
            [1., 1., 1., 1.]],
           [[1., 1., 1., 1.],
            [1., 1., 1., 1.]],
           [[1., 1., 1., 1.],
            [1., 1., 1., 1.]]])

b: tensor([[[0., 0., 0., 0.]],
           [[0., 0., 0., 0.]],
           [[0., 0., 0., 0.]]])

c: tensor([[[1., 1., 1., 1.],
            [1., 1., 1., 1.]],
           [[1., 1., 1., 1.],
            [1., 1., 1., 1.]],
           [[1., 1., 1., 1.],
            [1., 1., 1., 1.]]])

2、squeeze()

作用：与unsqueeze()相反，squeeze的作用是删除tensor中的所有维度为1的维度，注意只是删除维度，对数据没有影响。

举例说明：

import torch
a = torch.ones(3, 1, 4)   # 创建3*1*4的全为1的tensor
b = a.squeeze() # 删除b的维度为1的维度，即将b的shape转换为（3，4）
print('a:', a)
print('shape of a:', a.shape)
print('b:', b)
print('shape of b:', b.shape)

'''   运行结果   '''
a: tensor([[[1., 1., 1., 1.]],
           [[1., 1., 1., 1.]],
           [[1., 1., 1., 1.]]])
shape of a: torch.Size([3, 1, 4])

b: tensor([[1., 1., 1., 1.],
           [1., 1., 1., 1.],
           [1., 1., 1., 1.]])
shape of b: torch.Size([3, 4])