TF-GNN是基于TensorFlow的图神经网络框架,包含了图卷积网络(GCN)和图注意力网络(GAT)等常见的深度图神经网络。框架内部使用了稀疏矩阵等解决了图神经网络的计算效率和规模问题。
项目内容:
图神经网络(如GCN、GAT等)的工程化一直是一个大问题,在数据预处理、网络构建、性能优化等方面都需要细致的考虑。TF-GNN基于TensorFLow,为图神经网络提供了下面的功能:
图神经网络数据集的构建,包括将图神经网络存储格式、图数据的读取、训练测试划分等
基于稀疏矩阵的数据结构大大提升了框架的性能,在CPU上也可以轻松地运行图神经网络
提供了GCN和GAT层,可以直接构建GCN和GAT
数据集:
内置了下面的已经格式化的数据集:
DBLP(论文引用数据集)
M10(论文引用数据集)
示例:
图卷积网络示例(GraphConvolutionalNetworks):
# coding=utf-8from gnn.data.dataset import GraphDataset, WhiteSpaceTokenizerfrom gnn.data.example import load_M10, load_cora, load_dblpfrom gnn.model.gcn import GCN, GCNTrainerimport tensorflow as tf# eager mode must be enabledfrom tensorflow.contrib.eager.python import tfeimport osos.environ["CUDA_VISIBLE_DEVICES"] = "0"tfe.enable_eager_execution()# read graph dataset: data/M10 data/dblp# dataset = GraphDataset("data/dblp", ignore_featureless_node=True)dataset = load_M10("data/M10", ignore_featureless_node=True)adj = dataset.adj_matrix(sparse=True)feature_matrix, feature_masks = dataset.feature_matrix(bag_of_words=True, sparse=True)labels, label_masks = dataset.label_list_or_matrix(one_hot=False)train_node_indices, test_node_indices, train_masks, test_masks = dataset.split_train_and_test(training_rate=0.3)gcn_model = GCN([16, dataset.num_classes()])gcn_trainer = GCNTrainer(gcn_model)gcn_trainer.train(adj, feature_matrix, labels, train_masks, test_masks, learning_rate=1e-3, l2_coe=1e-3)
评论