torch.unsqueeze()函数理解

路过不知道

已于 2023-02-12 21:39:09 修改

阅读量4.4k

点赞数 5

分类专栏：深度学习文章标签： python 深度学习 pytorch

于 2023-02-12 17:29:18 首次发布

本文链接：https://blog.csdn.net/weixin_53142585/article/details/128997156

版权

深度学习专栏收录该内容

7 篇文章

订阅专栏

torch.unsqueeze()函数用于在给定的维度上增加一个新的尺寸，例如将一维张量转换为二维，二维转三维等。dim参数指定插入新维度的位置，可以是负数表示从末尾计数。文章通过示例展示了不同维度张量使用unsqueeze后的变化。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch.unsqueeze()函数理解

torch.unsqueeze(input, dim)  使用时等同于  input.unsqueeze(dim)

torch.unsqueeze()函数起到升维的作用，dim等于几表示在第几维度加一，比如原来x的size=([4])，x.unsqueeze(0)之后就变成了size=([1, 4])，而x.unsqueeze(1)之后就变成了size=([4, 1])，注意dim∈[-input.dim() - 1, input.dim() + 1]
例如

输入一维张量，即input.dim()=1

# 输入：
x = torch.tensor([1, 2, 3, 4])  # x.dim()=1
print(x)
print(x.shape)
y = x.unsqueeze(0)
print(y)
print(y.shape)  # 此时y.dim()=2
z = x.unsqueeze(1)
print(z)
print(z.shape)  # 此时z.dim()=2

# 输出：
tensor([1, 2, 3, 4])
torch.Size([4])
tensor([[1, 2, 3, 4]])
torch.Size([1, 4])
tensor([[1],
        [2],
        [3],
        [4]])
torch.Size([4, 1])

输入二维张量，即input.dim()=2

# 输入：
x = torch.tensor([[1, 2, 3], [4, 5, 6]])  # x.dim()=2
print(x)
print(x.shape)
y = x.unsqueeze(0)
print(y)
print(y.shape)  # 此时y.dim()=3
z = x.unsqueeze(1)
print(z)
print(z.shape)  # 此时z.dim()=3

# 输出：
tensor([[1, 2, 3],
        [4, 5, 6]])
torch.Size([2, 3])
tensor([[[1, 2, 3],
         [4, 5, 6]]])
torch.Size([1, 2, 3])
tensor([[[1, 2, 3]],

        [[4, 5, 6]]])
torch.Size([2, 1, 3])

输入四维张量，即input.dim()=4

# 输入：
x = torch.tensor([[[[1, 2, 3], 
                    [4, 5, 6]],
                [[0, 2, 1], 
                 [1, 5, 2]]],
                  
                [[[1, 2, 3], 
                  [4, 5, 6]],
                [[0, 2, 1], 
                 [1, 5, 2]]]])
print(x)
print(x.shape)

y2 = x.unsqueeze(2)
print(y2)
print(y2.shape)

y3 = x.unsqueeze(3)
print(y3)
print(y3.shape)

# 输出：
tensor([[[[1, 2, 3],
          [4, 5, 6]],

         [[0, 2, 1],
          [1, 5, 2]]],


        [[[1, 2, 3],
          [4, 5, 6]],

         [[0, 2, 1],
          [1, 5, 2]]]])
torch.Size([2, 2, 2, 3])
tensor([[[[[1, 2, 3],
           [4, 5, 6]]],


         [[[0, 2, 1],
           [1, 5, 2]]]],



        [[[[1, 2, 3],
           [4, 5, 6]]],


         [[[0, 2, 1],
           [1, 5, 2]]]]])
torch.Size([2, 2, 1, 2, 3])
tensor([[[[[1, 2, 3]],

          [[4, 5, 6]]],


         [[[0, 2, 1]],

          [[1, 5, 2]]]],



        [[[[1, 2, 3]],

          [[4, 5, 6]]],


         [[[0, 2, 1]],

          [[1, 5, 2]]]]])
torch.Size([2, 2, 2, 1, 3])