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

我该如何使用torch函数来实现这个呢？

相关问题：

• 暂无相关问题

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)

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

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

这个表示的是一个三维的数据结构，里面有10个部分，每个部分有1000个数据点，每个数据点有6个特征。

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

这个也是一个四维的数据结构，里面同样有10个部分，每个部分有1000个数据点，每个数据点有6个特征，不过这里多了一个维度，表示每个数据点还有一个额外的特征，虽然这个特征的值是1。