如何在 PyTorch 中将张量逐行乘以向量？

Question

当我有一个形状为 [12, 10] 的张量 m 和一个形状为 [12] 的标量向量 s 时，如何将 [=10] 的每一行相乘=] 与 s?

中相应的标量

Answer 1

您需要添加对应的单例维度：

m * s[:, None]

当 (12, 10) 张量乘以 (12, 1) 张量时，

s[:, None] 的大小为 (12, 1) pytorch 知道 broadcast s 沿第二个单例维度并正确执行“逐元素”乘积。

Answer 2

如果您事先知道维数并且可以硬编码 None 的正确数量，Shai 的答案就有效。这可以扩展到需要额外的维度：

mask = (torch.rand(12) > 0.5).int()  
data = (torch.rand(12, 2, 3, 4))

result = data * mask[:,None,None,None]

result.shape                  # torch.Size([12, 2, 3, 4])
mask[:,None,None,None].shape  # torch.Size([12, 1, 1, 1])

如果您正在处理可变或未知维度的数据，则可能需要手动将 mask 扩展到正确的形状

mask = (torch.rand(12) > 0.5).int()
while mask.dim() < data.dim(): mask.unsqueeze_(1)
result = data * mask

result.shape  # torch.Size([12, 2, 3, 4])
mask.shape    # torch.Size([12, 1, 1, 1])

这是一个有点丑陋的解决方案，但它确实有效。可能有一种更优雅的方法可以正确地重塑 mask 张量内联的可变维数

Answer 3

您可以将向量广播到更高维的张量like so:

def row_mult(input, vector):
    extra_dims = (1,)*(input.dim()-1)
    return t * vector.view(-1, *extra_dims)

如何在 PyTorch 中将张量逐行乘以向量？

How to multiply a tensor row-wise by a vector in PyTorch?

python

pytorch

tensor

scalar