python - PyTorch:如何通过广播两个不同形状的张量相乘

Question

我有以下两个 PyTorch 张量 A 和 B。

A = torch.tensor(np.array([40, 42, 38]), dtype = torch.float64)

tensor([40., 42., 38.], dtype=torch.float64)

B = torch.tensor(np.array([[[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5],[1,2,3,4,5]], [[4,5,6,7,8],[4,5,6,7,8],[4,5,6,7,8],[4,5,6,7,8],[4,5,6,7,8]], [[7,8,9,10,11],[7,8,9,10,11],[7,8,9,10,11],[7,8,9,10,11],[7,8,9,10,11]]]), dtype = torch.float64)

tensor([[[ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.]],

        [[ 4.,  5.,  6.,  7.,  8.],
         [ 4.,  5.,  6.,  7.,  8.],
         [ 4.,  5.,  6.,  7.,  8.],
         [ 4.,  5.,  6.,  7.,  8.],
         [ 4.,  5.,  6.,  7.,  8.]],

        [[ 7.,  8.,  9., 10., 11.],
         [ 7.,  8.,  9., 10., 11.],
         [ 7.,  8.,  9., 10., 11.],
         [ 7.,  8.,  9., 10., 11.],
         [ 7.,  8.,  9., 10., 11.]]], dtype=torch.float64)

张量 A 的形状为:

torch.Size([3])

张量 B 的形状为:

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

如何以这种方式将张量 A 与张量 B(使用广播)相乘，例如。张量 A 中的第一个值(即 40.)乘以张量 B 中第一个“嵌套”张量中的所有值，即。

tensor([[[ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.],
         [ 1.,  2.,  3.,  4.,  5.]],

分别对张量 A 中的其他 2 个值和张量 B 中的其他两个嵌套张量依此类推。

如果 A 和 B 都是形状为 (3,) 的数组，我可以用 numpy 数组做这个乘法(通过广播)——即。 A*B - 但我似乎无法用 PyTorch 张量找出与此对应的东西。非常感谢任何帮助。

Answer 1

在 pytorch(以及 numpy)中应用广播时，您需要从最后维度开始(查看https://pytorch.org/docs/stable/notes/broadcasting.html)。如果它们不匹配，则需要 reshape 张量。在您的情况下，它们不能直接广播:

      [3]  # the two values in the last dimensions are not one and do not match
[3, 5, 5]

相反，您可以在乘积之前重新定义 A = A[:, None, None] 以获得形状

[3, 1, 1]
[3, 5, 5]

满足广播条件