torch-geometric:torch_cluster库knn_graph函数详解

参考资料：https://github.com/zxin1023/pytorch_cluster

  今天在阅读一份图神经网络的代码时碰到了knn_graph函数。图神经网络如果使用动态图，通常需要获取节点的k近邻作为邻居节点，这个时候就会用到knn_graph函数了。knn_graph函数并不是torch-geometric自带的，而是torch_cluster库中的函数，这也可以说明为什么安装torch-geometric前需要先安装torch_cluster库。

  这里，我从官方手册里搬运knn_graph函数的参数说明和例子进行介绍：

  从这份官方文档中可以看到，这个函数是用来计算k近邻的，常用的参数有：
  （1）x：输入张量，即所有节点的特征，形状为[N, F]；
  （2）batch：一般每个epoch都会对一个batch的样本数据进行处理，而不是单个样本。这个参数是用来指代输入的N个节点分别属于batch中的哪个样本，参数形状为[N]。d