[发明专利]一种基于图神经网络知识蒸馏的图节点分类方法及装置

说明书

本发明实施例提供了一种基于图神经网络知识蒸馏的图节点分类方法及装置，将所述学生模型得到的第二预测结果，拟合所述教师模型得到的第一预测结果，通过使用教师模型的第一预测结果优化学生模型的，不需要教师模型和学生模型之间的集成或迭代，也能够提高分类精度，简化了优化学生模型的过程；学生模型采用标签传播公式，通过有标签的节点传播到相邻无标签的节点，可以从图结构的先验知识中受益；也通过特征变换公式，预测无标签点集的软标签，可以从特征的先验知识中受益，即数据集中具有硬标签的点集及无标签点集的特征。从而充分的使用先验知识，提高学生模型的分类效果，从而提高分类精度。

技术领域

本发明涉及机器学习技术领域，特别是涉及一种基于图神经网络知识蒸馏的图节点分类方法及装置。

背景技术

知识蒸馏被提出用于模型压缩，其中训练一个小的轻量级学生模型来模仿预训练的大教师模型的软预测。教师模型中的知识经过蒸馏以后，会转化为学生模型中的知识。

目前，知识蒸馏与图卷积网络(Graph convolutional networks，简称GCN)结合在一起进行应用。比如，可靠数据蒸馏(Resilient Distributed Dataset，简称RDD)用相同的体系结构，训练了多个GCN学生，然后对这些GCN学生进行集成以获得更好的性能。这样两个接收大小不同的GCN相互学习。其中，教师模型和学生模型，都是GCN。

但是，目前使用GCN作为学生模型的分类效果较差。

发明内容

本发明实施例的目的在于提供一种基于图神经网络知识蒸馏的图节点分类方法及装置，用以解决现有技术中使用GCN作为学生模型的分类效果较差的技术问题。具体技术方案如下：

第一方面，本发明实施例提供了基于图神经网络知识蒸馏的图节点分类方法，包括：

获取图神经网络教师模型，数据集，以及图结构中的节点集和边集，其中，所述教师模型为图神经网络GNN分类器，并且通过所述图结构及所述数据集预训练得到的，所述教师模型用于预测无标签点集的软标签可能为所述硬标签的第一概率，得到所述无标签点集的软标签的第一预测结果；所述数据集包括：具有硬标签的点集及无标签点集，所述每个节点表示数据集中的一个点；所述点为基准引文或购买商品；

基于所述第一预测结果，所述数据集及所述图结构，采用学生模型预测无标签点集的软标签，可能为所述硬标签的第二概率，得到所述无标签点集的软标签的第二预测结果；所述学生模型是基于所有硬标签的分布，采用标签传播公式，预测有标签的节点传播到相邻无标签的节点的概率分布，以及基于每个节点的特征，采用特征变换公式，预测无标签点集的软标签进行组合训练的；

将所述学生模型得到的第二预测结果，拟合所述教师模型得到的第一预测结果，直至所述学生模型得到的第二预测结果与所述教师模型得到的第一预测结果最相近，得到训练好的学生模型；所述训练好的学生模型用于对无标签点集的软标签的预测，得到基准引文的类型或购买商品的类型。

进一步的，在所述基于所述第一预测结果，所述数据集及所述图结构，采用学生模型预测无标签点集的软标签，可能为所述硬标签的第二概率，得到所述无标签点集的软标签的第二预测结果之前，所述方法还包括：

根据如下公式：

为每个节点初始化标签预测，以完成初始化学生模型；

其中，为节点v在第k次迭代中的预测概率分布，为节点v在初始化的预测概率分布，LP为标签传播的英文简称，∈为属于，为实数集合，Y为标签集合，|.|为集合的基数，为任意，V_L为有标签节点集合，V为所有节点，L为有标签，V_U为无标签节点集合，U为无标签，f_LP为所述第一预测结果中的LP的最终预测。