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

【问题标题】：How to vectorise a list of matrix vector multiplications using pytorch/numpy如何使用 pytorch/numpy 向量化矩阵向量乘法列表
【发布时间】：2018-12-02 18:14:05
【问题描述】：

例如，我有一个NB 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 操作）

【问题讨论】：

你能加入minimal reproducible example吗？
第二。我对您正在尝试的内容有一个想法，但如果我有一些定义的输入和输出要测试，那会更令人放心。

标签： python numpy matrix-multiplication pytorch tensor

【解决方案1】：

您可以使用torch.bmm() 和一些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]])

【讨论】：