使用 pytorch 张量进行维度扩展

Question

pytorch张量的维度扩展方式是什么？

-之前： torch.Size([3, 3, 3])

tensor([[[ 0.,  1.,  2.],
         [ 3.,  4.,  5.],
         [ 6.,  7.,  8.]],

        [[ 9., 10., 11.],
         [12., 13., 14.],
         [15., 16., 17.]],

        [[18., 19., 20.],
         [21., 22., 23.],
         [24., 25., 26.]]], device='cuda:0', dtype=torch.float64)

-之后： torch.Size([2, 3, 3, 3])

tensor([[[[ 0.,  1.,  2.],
          [ 3.,  4.,  5.],
          [ 6.,  7.,  8.]],

         [[ 9., 10., 11.],
          [12., 13., 14.],
          [15., 16., 17.]],

         [[18., 19., 20.],
          [21., 22., 23.],
          [24., 25., 26.]]],


        [[[0., 1., 2.],
          [ 3.,  4.,  5.],
          [ 6.,  7.,  8.]],

         [[ 9., 10., 11.],
          [12., 13., 14.],
          [15., 16., 17.]],

         [[18., 19., 20.],
          [21., 22., 23.],
          [24., 25., 26.]]]], device='cuda:0', dtype=torch.float64)

在 numpy 下会像这样工作：

b =  np.broadcast_to(a1[None, :,:,:], (2,3,3,3))

这个在pytorch下如何工作？我想利用 gpu。在此先感谢您的帮助！

Answer 1

可以使用unsqeeze添加一个新的维度（下面使用0来指定第一个维度，即位置0），然后沿着该维度重复数据两次（一次，即没有重复，沿着其他维度）。

before = torch.tensor(..., dtype=torch.float64, device='cuda:0')
after = before.unsqueeze(0).repeat(2, 1, 1, 1)

Answer 2

我们可以使用 torch.Tensor.expand 来获得您给定的预期结果

b = a1.expand([2, 3, 3, 3])

使用 pytorch 张量进行维度扩展

dimension extension with pytorch tensors

pytorch

gpu

tensor

dimensions