python-3.x - 将许多矩阵和许多向量相乘 pytorch

Question

我正在尝试将以下内容相乘：

一批矩阵N x M x D
一批向量 N x D x 1
得到一个结果：N x M x 1

好像我在做N点积M x D D x 1。

我似乎无法在 PyTorch 中找到正确的功能。

torch.bmm据我所知，仅适用于一批向量和一个矩阵。如果我必须使用，torch.einsum那就这样吧，但我宁愿不要！

标签： python-3.xmatrixpytorchlinear-algebra

Answer 1

它非常简单直观einsum：

torch.einsum('ijk, ikl->ijl', mats, vecs)

但你的操作只是：

mats @ vecs