在节点特征的某些子集上构建 KNN 图

【问题标题】：Build KNN graph over some subset of Node features在节点特征的某些子集上构建 KNN 图
【发布时间】：2021-02-24 05:49:24
【问题描述】：

我有一个点云，我想在其上使用图神经网络。点云中的每个点都以其位置坐标和颜色为特征。所以单个节点是(X, Y, Z, C)。

现在我想对此应用边缘卷积（如 in the DGL Edge-Conv example 所述，为此我应该在 (X, Y, Z)（而不是 C）上构建最近邻图，然后使用所有 4 个属性作为我的神经网络的特征。

什么是干净有效的方法来做到这一点？（我有很多数据，所以我想很好地批处理和整理）

【问题讨论】：

标签： machine-learning deep-learning pytorch dgl

【解决方案1】：

假设您有一个形状为(NUM_POINTS, 4) 的张量pc，其中每一行都是(X, Y, Z, C)，那么您可以使用sklearn，如下所示：

from sklearn.neighbors import NearestNeighbors
import dgl

k = 3 # number of neighbours you want

neigh = NearestNeighbors(n_neighbors=k)
neigh.fit(pc[:, :3].numpy()) # selects only (X, Y, Z)

knn = neigh.kneighbors_graph()
graph = dgl.from_scipy(knn)
graph.ndata['x'] = pc

我建议将这些图表保存到磁盘，这样就不会在每次训练等时计算它们。

【讨论】：