Question

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

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

Answer 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]])

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

1 个答案: