matrix - 在 Pytorch 中创建 knn 相邻矩阵

Question

在 Pytorch 中，假设我有一个 top-k 索引矩阵 P(B,N,k) , 权重矩阵 W(B,N,N)和一个目标矩阵 A(B,N,N) ，我想得到一个作为以下循环运行的相邻矩阵:

for i in range(B):
     for ii in range(N):
         for j in range(k):
             if weighted:
                 A[i][ii][P[i][ii][j]] = W[i][ii][P[i][ii][j]]
             else:
                 A[i][ii][P[i][ii][j]] = 1

如何在 中更高效、更简洁地实现它pytorch ?

Answer 1

我认为您正在寻找 torch.scatter_ :

A.scatter_(dim=2, index=P, src=W)  # for the weighted version
A.scatter_(dim=2, index=P, src=torch.ones_like(W))  # for the un-weighted version