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

Question

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

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

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

Answer 1

假设你有一个pc形状张量(NUM_POINTS, 4)，每行是(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

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