如何使用 pytorch/numpy 向量化矩阵向量乘法列表

Question

例如，我有一个 N B x H 张量列表（即 N x B x H 张量）和一个 N 向量列表（即 N x B张量）。我想将列表中的每个 B x H 张量与相应的 B 维张量相乘，得到一个 N x H 张量。

我知道如何使用单个for-loop 和PyTorch 来实现计算，但是有没有矢量化植入？ （即没有 for-loop，仅使用 PyTorch/numpy 操作）

Answer 1

您可以通过 torch.bmm() and some torch.squeeze()/torch.unsqueeze() 实现。

我个人比较喜欢更通用的torch.einsum()（我觉得它更具可读性）：

import torch
import numpy as np

A = torch.from_numpy(np.array([[[1, 10, 100], [2, 20, 200], [3, 30, 300]],
                               [[4, 40, 400], [5, 50, 500], [6, 60, 600]]]))
B = torch.from_numpy(np.array([[ 1,  2,  3],
                               [-1, -2, -3]]))

AB = torch.einsum("nbh,nb->nh", (A, B))
print(AB)
# tensor([[   14,   140,  1400],
#         [  -32,  -320, -3200]])