我可以得到以下结果没有循环，只是使用矩阵运算

Question

我想得到以下结果，如输出显示，我可以使它没有for循环，只是使用矩阵操作，如numpy或torch，谢谢大家。

import torch
torch.manual_seed(0)
input = torch.rand((1,3543,768))
w = torch.rand((3543,768,1))
output = torch.zeros(input.shape[0], input.shape[1], 1)
for i in range(input.shape[1]):
output[:, i, :] = torch.matmul(input[:,i,:], w[i,:,:])

输出:

tensor([[[0.4963, 0.7682, 0.0885,  ..., 0.1497, 0.3923, 0.9338],
[0.1164, 0.3539, 0.6640,  ..., 0.5025, 0.4458, 0.2083],
[0.3337, 0.6611, 0.7212,  ..., 0.0513, 0.3007, 0.8048],
...,
[0.5168, 0.3856, 0.0275,  ..., 0.6173, 0.8562, 0.8301],
[0.8208, 0.1885, 0.4760,  ..., 0.4657, 0.4911, 0.2394],
[0.8758, 0.9653, 0.8739,  ..., 0.7802, 0.1038, 0.4542]]])

Answer 1

我已经解决了

output = torch.einsum("bij,ijk->bik", input, w)