PyTorch。逐行点阵产品

Question

假设我有两个时序。

a = torch.randn(10, 1000, 1, 4)
b = torch.randn(10, 1000, 6, 4)

其中第三个指数是一个向量的指数。

我想取每个向量之间的点积，在 b 的向量，在 a.

为了说明，这就是我的意思。

dots = torch.Tensor(10, 1000, 6, 1)
for b in range(10):
     for c in range(1000):
           for v in range(6):
            dots[b,c,v] = torch.dot(b[b,c,v], a[b,c,0]) 

我如何用火炬函数来实现这个目标?

Answer 1

a = torch.randn(10, 1000, 1, 4)
b = torch.randn(10, 1000, 6, 4)

c = torch.sum(a * b, dim=-1)

print(c.shape)

torch.Size([10, 1000, 6])

c = c.unsqueeze(-1)
print(c.shape)

torch.Size([10, 1000, 6, 1])