【Pytorch】unsqueeze() 方法

---

一、unsqueeze() 方法

unsqueeze() 方法起升维的作用，参数表示在哪个地方加一个维度。

---

二、实例

import torch
input=torch.arange(0,6)
print(input)
print(input.shape)

print(input.unsqueeze(0))    # 在0维处加一维
print(input.unsqueeze(0).shape)

print(input.unsqueeze(1))    # 在1维处加一维
print(input.unsqueeze(1).shape)

输出结果：

tensor([0, 1, 2, 3, 4, 5])
torch.Size([6])
tensor([[0, 1, 2, 3, 4, 5]])
torch.Size([1, 6])  # 在0维处加了一维
tensor([[0],
        [1],
        [2],
        [3],
        [4],
        [5]])
torch.Size([6, 1])  # 在1维处加了一维

---

参考链接

1. pytorch中unsqueeze()函数理解