pytorch-unsqueeze用法

tensor.unsqueeze 为tenor添加维度

a = torch.rand((1, 3, 5), dtype=torch.float32)
print("a shape is ", a.size())

b = a.unsqueeze(-1)  # 负1表示 在最后一维上添加
print("b shape is ", b.size())

输出结果如下

a shape is  torch.Size([1, 3, 5])
b shape is  torch.Size([1, 3, 5, 1])

真实数值如下：

a:
tensor([[[0.6146, 0.2028, 0.4266, 0.9713, 0.1965],
         [0.3768, 0.5539, 0.4202, 0.5306, 0.3155],
         [0.1831, 0.1981, 0.6740, 0.1039, 0.2108]]])

b:
tensor([[[[0.6146],
          [0.2028],
          [0.4266],
          [0.9713],
          [0.1965]],

         [[0.3768],
          [0.5539],
          [0.4202],
          [0.5306],
          [0.3155]],

         [[0.1831],
          [0.1981],
          [0.6740],
          [0.1039],
          [0.2108]]]])

注意：维度的顺序是从外往里数的， 看a和b就能看看出来。