【Pytorch】unsqueeze() 方法

文章目录
一、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维处加了一维
 
参考链接
pytorch中unsqueeze()函数理解