Pytorch 批量矩阵向量行乘法

Pytorch batchwise matrix vector rowwise multiplication

我正在使用 pytorch。

如果我有一个大小为 (d1, d2) 的矩阵 M 和一个大小为 d2 的向量 V，执行 M*V 会得到一个输出 OUT 的大小 (d1, d2)，其中 M 的每一行都乘以 V。

我需要分批做同样的事情，其中​​矩阵 M 是固定的，我有一批 dB 个向量。

在实践中，给定大小 (d1, d2) 的张量 M 和大小 (dB, d2) 的张量 V，我需要得到输出 OUT大小为 (dB, d1, d2) 的张量，因此 OUT[i] = M*V[i].

如何使用pytorch高效获取？

您可以使用 Einstein Notation 来实现：

torch.einsum('ij,bj->bij', M, V)

这个简单的技巧可以解决这个问题：

M.unsqueeze(0) * V.unsqueeze(1)

这会对形状为 (1, d1, d2) 和 (dB, 1, d2) 的张量进行乘法运算，得到形状为 (dB, d1, d2).

的所需输出