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)

炬。尺寸([10， 1000， 6](

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

炬。尺寸([10， 1000， 6， 1](