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）。如果它们不匹配，您需要重塑张量。在您的情况下，它们不能直接广播：

      [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]

满足广播条件。