pytorch函数unsqueeze函数的理解

在看许多pytorch的代码时，为了计算上的方便，通常会用到unsqueeze函数，一直不得要领，这次专门去做个实验学习一下。官方文档对这个函数描述如下，就是在指定的位置插入一个维度，有两个参数，input是输入的tensor,dim是要插到的维度需要注意的是dim的范围是[-input.dim()-1, input.dim()+1)，是一个左闭右开的区间，当dim为负值时，会自动转换为dim =

ljwwjl

42953人浏览 · 2021-03-31 09:49:26

ljwwjl · 2021-03-31 09:49:26 发布

在看许多pytorch的代码时，为了计算上的方便，通常会用到unsqueeze函数，一直不得要领，这次专门去做个实验学习一下。

官方文档对这个函数描述如下，就是在指定的位置插入一个维度，有两个参数，input是输入的tensor,dim是要插到的维度

需要注意的是dim的范围是[-input.dim()-1, input.dim()+1)，是一个左闭右开的区间，当dim为负值时，会自动转换为dim = dim+input.dim()+1，类似于使用负数对python列表进行切片。

下面使用一个二维矩阵看下dim不同时呈现出的效果：

# 创建一个3*4的全1二维tensor
a = torch.ones(3,4)
'''
运行结果
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
'''

在0维度上插入一个维度，可以看到现在a的形状变为[1, 3, 4]，第0维度的大小默认是1

a = a.unsqueeze(0)
print(a.shape)
'''
运行结果
tensor([[[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]]])

torch.Size([1, 3, 4])
'''

在最后一个维度上插入一个维度，形状变为[3, 4, 1]

a = a.unsqueeze(a.dim())
print(a.shape)
'''
运行结果
tensor([[[1.],
         [1.],
         [1.],
         [1.]],

        [[1.],
         [1.],
         [1.],
         [1.]],

        [[1.],
         [1.],
         [1.],
         [1.]]])

torch.Size([3, 4, 1])
'''