在节点要素的某些子集上构建KNN图

问题描述

我有一个点云，我想在其上使用图神经网络。点云中的每个点都由其位置坐标及其颜色来表征。因此，一个节点是(X,Y,Z,C)。

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

什么是清洁有效的方法？ （我有很多数据，所以我想很好地进行批处理和整理）

解决方法

假设您有一个形状为 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,Z)

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

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