pytorch - 张量中仅一维的矩阵向量乘法

Question

是否可以仅将张量中的一个（最后一个）维度与其他向量相乘？

例如，假设一个张量 T=[100, 20, 400] 和一个矩阵 M =[400, 400]。是否可以进行操作h_{transpose}*M*h，h张量 T 中的最后一个维度在哪里？换句话说，是否可以使用（可能是 pytorch）内置函数来获得大小为 [100, 20, 1] 的结果张量？

标签： pytorch

Answer 1

我认为最简单（当然是最短）的解决方案是使用einsum.

import torch

T = torch.randn(100, 20, 400)
M = torch.randn(400, 400)

res = torch.einsum('abc,cd,abd->ab', (T, M, T)).unsqueeze(-1)

它基本上说“对于所有(a, b, c, d)有界的，乘以并T[a, b, c]累积它”。M[c, d]T[a, b, d]res[a, b]

由于einsum是根据 等基本构建块实现的mm，transpose因此这当然可以展开为更“经典”的解决方案，但现在我的大脑对此感到失望。