pytorch中的squeeze和unsqueeze

pytorch中的squeeze和unsqueeze

unsqueeze即在参数指定的维度位置，增加一个维度（就是在第几个“[”的位置增加一个“[”）

import torch

a = torch.arange(0,8)
print(a)
b = a.view(2,4)
print(b)
b = b.unsqueeze(1)
print(b)

tensor([0, 1, 2, 3, 4, 5, 6, 7])
tensor([[0, 1, 2, 3],
        [4, 5, 6, 7]])
tensor([[[0, 1, 2, 3]],

        [[4, 5, 6, 7]]])

```python
import torch

a = torch.arange(0,8)
print(a)
b = a.view(2,4)
print(b)
b = b.unsqueeze(0)
print(b)

tensor([0, 1, 2, 3, 4, 5, 6, 7])
tensor([[0, 1, 2, 3],
        [4, 5, 6, 7]])
tensor([[[0, 1, 2, 3],
         [4, 5, 6, 7]]])

squeeze即去除一个维度（这个维度只能为1）

import torch

a = torch.arange(0,8)
print(a)
b = a.view(1,2,4)
print(f"b's shape is {b.shape} \n {b}")
b = b.squeeze(-3)
print(f"b's shape is {b.shape} \n {b}")

tensor([0, 1, 2, 3, 4, 5, 6, 7])
b's shape is torch.Size([1, 2, 4]) 
 tensor([[[0, 1, 2, 3],
         [4, 5, 6, 7]]])
b's shape is torch.Size([2, 4]) 
 tensor([[0, 1, 2, 3],
        [4, 5, 6, 7]])