WikiLinks GNN 链路预测

基于图神经网络(GNN)的维基百科页面间链接预测与知识图谱补全

Python 3.8+
PyTorch 2.0+
AUC: 0.798
MIT License

项目介绍

WikiLinks GNN链路预测项目是一个基于图神经网络(GNN)的机器学习项目,旨在预测维基百科页面之间的潜在链接关系。本项目通过分析页面间的结构特征,构建了一个高效的链路预测模型,可用于知识图谱补全、推荐系统等应用场景。

核心问题

在维基百科的庞大知识网络中,许多页面之间的关联关系尚未被明确记录。本项目通过机器学习方法,自动识别这些潜在的链接关系,从而:

项目亮点

先进的图神经网络

使用PyTorch Geometric实现多种GNN架构,包括GCN、GAT和GraphSAGE,支持复杂的图结构学习。

全面的实验评估

提供完整的实验流程,包括数据预处理、模型训练、超参数调优和结果可视化,确保结果可复现。

模块化代码设计

代码结构清晰,模块化设计,便于扩展和二次开发,适合学术研究和工业应用。

开箱即用

提供预训练模型、一键训练脚本和Google Colab演示,降低使用门槛,快速上手。

实验结果

性能指标

0.798
测试集 AUC
0.784
测试集 AP
0.763
F1分数
0.696
准确率

模型配置

  • 总参数量: 21,729 (全部可训练)
  • 训练数据: 420条边
  • 验证数据: 140条边
  • 测试数据: 140条边
  • 训练轮次: 300 epochs
  • 最佳验证AUC: 0.897 (Epoch 110)

技术栈

  • 深度学习框架: PyTorch 2.0+
  • 图神经网络库: PyTorch Geometric
  • 数据处理: NumPy, Pandas
  • 可视化: Matplotlib, Seaborn
  • 实验管理: JSON日志, TensorBoard(可选)

支持模型

  • 图卷积网络 (GCN)
  • 图注意力网络 (GAT)
  • GraphSAGE
  • 链路预测专用解码器

训练过程可视化

通过动态GIF展示模型在300个训练轮次中的学习过程,直观呈现损失下降和AUC提升的趋势。

WikiLinks GNN训练过程动画

图: 训练过程中损失函数下降和验证AUC提升的动态变化

在训练过程中,模型在第110个epoch达到最佳验证AUC(0.897),最终在测试集上获得0.798的AUC分数。训练损失持续下降表明模型有效学习到了数据中的模式。

快速开始

1. 克隆仓库

git clone https://github.com/inneedloveBu/wikinet-link-prediction.git
cd wikinet-link-prediction

2. 安装依赖

pip install -r requirements.txt

3. 训练模型

python train.py --model gcn --epochs 300 --lr 0.001

4. 使用预训练模型进行预测

from models.link_prediction import WikiLinksPredictor

# 加载预训练模型
predictor = WikiLinksPredictor.load_pretrained()

# 进行预测
predictions = predictor.predict(new_data)
print(f"预测结果: {predictions}")

参考文献

本项目基于以下重要的研究工作:

  1. Kipf, T. N., & Welling, M. (2016). Semi-Supervised Classification with Graph Convolutional Networks. arXiv preprint arXiv:1609.02907.
  2. Hamilton, W., Ying, Z., & Leskovec, J. (2017). Inductive Representation Learning on Large Graphs. Advances in neural information processing systems, 30.
  3. Zhang, M., & Chen, Y. (2018). Link Prediction Based on Graph Neural Networks. Advances in Neural Information Processing Systems.
如果本项目的代码或思路对您的研究有帮助,请考虑引用上述相关文献。