tf_geometric - TensorFlow 图神经网络框架
GPL
跨平台
Python
软件简介
tf_geometric是一个高效且友好的图神经网络框架，同时支持TensorFlow 1.x和2.x。
tf_geometric使用消息机制来实现图神经网络，这会比基于普通矩阵实现的版本更为高效且比基于稀疏矩阵实现的版本更为友好。另外，框架为复杂的图神经网络操作提供了简单优雅的API。下面的示例构建了一个图，并对其执行了多头图注意力网络（GAT）：
# coding=utf-8
import numpy as np
import tf_geometric as tfg
import tensorflow as tf
graph = tfg.Graph(
x=np.random.randn(5, 20), # 5 nodes, 20 features,
edge_index=[[0, 0, 1, 3],
[1, 2, 2, 1]] # 4 undirected edges
)
print("Graph Desc: \n", graph)
graph.convert_edge_to_directed() # pre-process edges
print("Processed Graph Desc: \n", graph)
print("Processed Edge Index:\n", graph.edge_index)
# Multi-head Graph Attention Network (GAT)
gat_layer = tfg.layers.GAT(units=4, num_heads=4, activation=tf.nn.relu)
output = gat_layer([graph.x, graph.edge_index])
print("Output of GAT: \n", output)
Output:
Graph Desc:
Graph Shape: x => (5, 20) edge_index => (2, 4) y => None
Processed Graph Desc:
Graph Shape: x => (5, 20) edge_index => (2, 8) y => None
Processed Edge Index:
[[0 0 1 1 1 2 2 3]
[1 2 0 2 3 0 1 1]]
Output of GAT:
tf.Tensor(
[[0.22443159 0. 0.58263206 0.32468423]
[0.29810357 0. 0.19403605 0.35630274]
[0.18071976 0. 0.58263206 0.32468423]
[0.36123228 0. 0.88897204 0.450244 ]
[0. 0. 0.8013462 0. ]], shape=(5, 4), dtype=float32)
面向对象和函数式API
框架同时提供了面向对象和函数式API，利用这些API你可以构建许多很酷的东西：
# coding=utf-8
import os
# Enable GPU 0
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tf_geometric as tfg
import tensorflow as tf
import numpy as np
from tf_geometric.utils.graph_utils import convert_edge_to_directed
# ==================================== Graph Data Structure ====================================
# In tf_geometric, graph data can be either individual Tensors or Graph objects
# A graph usually consists of x(node features), edge_index and edge_weight(optional)
# Node Features => (num_nodes, num_features)
x = np.random.randn(5, 20).astype(np.float32) # 5 nodes, 20 features
# Edge Index => (2, num_edges)
# Each column of edge_index (u, v) represents an directed edge from u to v.
# Note that it does not cover the edge from v to u. You should provide (v, u) to cover it.
# This is not convenient for users.
# Thus, we allow users to provide edge_index in undirected form and convert it later.
# That is, we can only provide (u, v) and convert it to (u, v) and (v, u) with `convert_edge_to_directed` method.
edge_index = np.array([
[0, 0, 1, 3],
[1, 2, 2, 1]
])
# Edge Weight => (num_edges)
edge_weight = np.array([0.9, 0.8, 0.1, 0.2]).astype(np.float32)
# Make the edge_index directed such that we can use it as the input of GCN
edge_index, edge_weight = convert_edge_to_directed(edge_index, edge_weight=edge_weight)
# We can convert these numpy array as TensorFlow Tensors and pass them to gnn functions
outputs = tfg.nn.gcn(
tf.Variable(x),
tf.constant(edge_index),
tf.constant(edge_weight),
tf.Variable(tf.random.truncated_normal([20, 2])) # GCN Weight
)
print(outputs)
# Usually, we use a graph object to manager these information
# edge_weight is optional, we can set it to None if you don't need it
graph = tfg.Graph(x=x, edge_index=edge_index, edge_weight=edge_weight)
# You can easily convert these numpy arrays as Tensors with the Graph Object API
graph.convert_data_to_tensor()
# Then, we can use them without too many manual conversion
outputs = tfg.nn.gcn(
graph.x,
graph.edge_index,
graph.edge_weight,
tf.Variable(tf.random.truncated_normal([20, 2])), # GCN Weight
cache=graph.cache # GCN use caches to avoid re-computing of the normed edge information
)
print(outputs)
# For algorithms that deal with batches of graphs, we can pack a batch of graph into a BatchGraph object
# Batch graph wrap a batch of graphs into a single graph, where each nodes has an unique index and a graph index.
# The node_graph_index is the index of the corresponding graph for each node in the batch.
# The edge_graph_index is the index of the corresponding edge for each node in the batch.
batch_graph = tfg.BatchGraph.from_graphs([graph, graph, graph, graph])
# Graph Pooling algorithms often rely on such batch data structure
# Most of them accept a BatchGraph's data as input and output a feature vector for each graph in the batch
outputs = tfg.nn.mean_pooling(batch_graph.x, batch_graph.node_graph_index, num_graphs=batch_graph.num_graphs)
print(outputs)
# We can reversely split a BatchGraph object into Graphs objects
graphs = batch_graph.to_graphs()
# ==================================== Built-in Datasets ====================================
# all graph data are in numpy format
train_data, valid_data, test_data = tfg.datasets.PPIDataset().load_data()
# we can convert them into tensorflow format
test_data = [graph.convert_data_to_tensor() for graph in test_data]
# ==================================== Basic OOP API ====================================
# OOP Style GCN (Graph Convolutional Network)
gcn_layer = tfg.layers.GCN(units=20, activation=tf.nn.relu)
for graph in test_data:
# Cache can speed-up GCN by caching the normed edge information
outputs = gcn_layer([graph.x, graph.edge_index, graph.edge_weight], cache=graph.cache)
print(outputs)
# OOP Style GAT (Multi-head Graph Attention Network)
gat_layer = tfg.layers.GAT(units=20, activation=tf.nn.relu, num_heads=4)
for graph in test_data:
outputs = gat_layer([graph.x, graph.edge_index])
print(outputs)
# ==================================== Basic Functional API ====================================
# Functional Style GCN
# Functional API is more flexible for advanced algorithms
# You can pass both data and parameters to functional APIs
gcn_w = tf.Variable(tf.random.truncated_normal([test_data[0].num_features, 20]))
for graph in test_data:
outputs = tfg.nn.gcn(graph.x, edge_index, edge_weight, gcn_w, activation=tf.nn.relu)
print(outputs)
# ==================================== Advanced OOP API ====================================
# All APIs are implemented with Map-Reduce Style
# This is a gcn without weight normalization and transformation.
# Create your own GNN Layer by subclassing the MapReduceGNN class
class NaiveGCN(tfg.layers.MapReduceGNN):
def map(self, repeated_x, neighbor_x, edge_weight=None):
return tfg.nn.identity_mapper(repeated_x, neighbor_x, edge_weight)
def reduce(self, neighbor_msg, node_index, num_nodes=None):
return tfg.nn.sum_reducer(neighbor_msg, node_index, num_nodes)
def update(self, x, reduced_neighbor_msg):
return tfg.nn.sum_updater(x, reduced_neighbor_msg)
naive_gcn = NaiveGCN()
for graph in test_data:
print(naive_gcn([graph.x, graph.edge_index, graph.edge_weight]))
# ==================================== Advanced Functional API ====================================
# All APIs are implemented with Map-Reduce Style
# This is a gcn without without weight normalization and transformation
# Just pass the mapper/reducer/updater functions to the Functional API
for graph in test_data:
outputs = tfg.nn.aggregate_neighbors(
x=graph.x,
edge_index=graph.edge_index,
edge_weight=graph.edge_weight,
mapper=tfg.nn.identity_mapper,
reducer=tfg.nn.sum_reducer,
updater=tfg.nn.sum_updater
)
print(outputs)
示例
- Graph Convolutional Network (GCN)
- Multi-head Graph Attention Network (GAT)