tf_geometric

高效且易用的图神经网络（GNN）库，适用于 TensorFlow 1.x 和 2.x。

受 rusty1s/pytorch_geometric 的启发，我们为 TensorFlow 构建了一个 GNN 库。

主页与文档

• 主页：https://github.com/CrawlScript/tf_geometric
• 文档：https://tf-geometric.readthedocs.io (中文版)
• 论文：使用 tf_geometric 在 TensorFlow 中实现高效的图深度学习

高效且友好

我们采用消息传递机制来实现图神经网络（GNN），这种方式比基于稠密矩阵的实现更高效，也比基于稀疏矩阵的实现更加友好。此外，我们还为复杂的 GNN 操作提供了简单而优雅的 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个节点，每个节点有20个特征，
    edge_index=[[0, 0, 1, 3],
                [1, 2, 2, 1]]  # 4条无向边
)

print("图描述：\n", graph)

graph = graph.to_directed()  # 处理边信息
print("处理后的图描述：\n", graph)
print("处理后的边索引：\n", graph.edge_index)

# 多头图注意力网络（GAT）
gat_layer = tfg.layers.GAT(units=4, num_heads=4, activation=tf.nn.relu)
output = gat_layer([graph.x, graph.edge_index])
print("GAT 的输出：\n", output)

输出：

图描述：
 图结构：x => (5, 20)	edge_index => (2, 4)	y => None

处理后的图描述：
 图结构：x => (5, 20)	edge_index => (2, 8)	y => None

处理后的边索引：
 [[0 0 1 1 1 2 2 3]
 [1 2 0 2 3 0 1 1]]

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)

DEMO

我们建议您从一些示例开始入手。

节点分类

• 图卷积网络（GCN）
• 多头图注意力网络（GAT）
• 近似个性化神经预测传播（APPNP）
• 大型图上的归纳表示学习（GraphSAGE）
• 具有快速局部化谱滤波的图卷积神经网络（ChebyNet）
• 简单图卷积（SGC）
• 拓扑自适应图卷积网络（TAGCN）
• 深度图信息最大化（DGI）
• DropEdge：面向节点分类的深层图卷积网络
• 用于文本分类的图卷积网络（TextGCN）
• 简单谱图卷积（SSGC/S^2GC）

图分类

• 均值池化（MeanPooling）
• 图同构网络（GIN）
• 自注意力图池化（SAGPooling）
• 可微分池化的层次化图表示学习（DiffPool）
• 顺序很重要：集合的序列到序列模型（Set2Set）
• ASAP：用于学习层次化图表示的自适应结构感知池化（ASAP）
• 面向图分类的端到端深度学习架构（SortPool）
• 基于图神经网络的谱聚类用于图池化（MinCutPool）

链接预测

• 图自编码器（GAE）

模型保存与加载

• 模型保存与加载
• 使用 tf.train.Checkpoint 保存和加载模型

分布式训练

• 面向节点分类的分布式 GCN
• 面向图分类的分布式均值池化

稀疏数据

• 稀疏节点特征

安装

要求：

• 操作系统：Windows / Linux / Mac OS
• Python：版本 >= 3.7
• Python 包：
  • tensorflow/tensorflow-gpu：>= 1.15.0 或 >= 2.7.0
  • tf_sparse
  • numpy >= 1.17.4
  • networkx >= 2.1
  • scipy >= 1.1.0

请使用以下命令之一进行安装：

pip install -U tf_geometric # 这不会安装 tensorflow/tensorflow-gpu 包

pip install -U tf_geometric[tf1-cpu] # 这将安装 TensorFlow 1.x CPU 版本

pip install -U tf_geometric[tf1-gpu] # 这将安装 TensorFlow 1.x GPU 版本

pip install -U tf_geometric[tf2-cpu] # 这将安装 TensorFlow 2.x CPU 版本

pip install -U tf_geometric[tf2-gpu] # 这将安装 TensorFlow 2.x GPU 版本

面向对象与函数式 API

我们同时提供了面向对象和函数式 API，您可以利用它们创造出许多有趣的东西。

# coding=utf-8
import os
# 启用 GPU 0
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import tf_geometric as tfg
import tensorflow as tf
import numpy as np

# ==================================== 图数据结构 ====================================
# 在 tf_geometric 中，图的数据可以用一组张量（numpy.ndarray 或 tf.Tensor）或者一个 tfg.Graph 对象来表示。
# 一个图通常由节点特征 x、边索引 edge_index 和边权重 edge_weight（可选）组成。

# 节点特征 => (num_nodes, num_features)
x = np.random.randn(5, 20).astype(np.float32)  # 5个节点，每个节点有20个特征

# 边索引 => (2, num_edges)
# 边索引的每一列 (u, v) 表示一条从 u 到 v 的有向边。
# 注意，它并不包含从 v 到 u 的边。你需要提供 (v, u) 来补充这一部分。
# 这种方式对用户不太友好。因此，我们允许用户以无向形式提供边索引，并在后续将其转换为有向形式。
# 也就是说，我们可以只提供 (u, v)，然后通过 `convert_edge_to_directed` 方法将其转换为 (u, v) 和 (v, u)。
edge_index = np.array([
    [0, 0, 1, 3],
    [1, 2, 2, 1]
])

# 边权重 => (num_edges)
edge_weight = np.array([0.9, 0.8, 0.1, 0.2]).astype(np.float32)


# 通常，我们会使用一个图对象来管理这些信息。
# 边权重是可选的，如果你不需要的话可以设置为 None。
# 使用 'to_directed' 方法将图转换为有向边形式，以便作为 GCN 的输入。
graph = tfg.Graph(x=x, edge_index=edge_index, edge_weight=edge_weight).to_directed()


# 定义一个图卷积层（GCN）
gcn_layer = tfg.layers.GCN(4, activation=tf.nn.relu)

# 在图上执行 GCN
h = gcn_layer([graph.x, graph.edge_index, graph.edge_weight])
print("节点表示（图上的 GCN）：\n", h)

for _ in range(10):
    # 使用 Graph.cache 可以避免重新计算 GCN 的归一化邻接矩阵，
    # 从而显著提高 GCN 的效率。
    h = gcn_layer([graph.x, graph.edge_index, graph.edge_weight], cache=graph.cache)


# 对于处理图批次的算法，我们可以将一批图打包成一个 BatchGraph 对象
# Batch graph 将一批图包装成一个单一的图，其中每个节点都有一个唯一的索引和一个图索引。
# node_graph_index 是批次中每个节点对应图的索引。
# edge_graph_index 是批次中每个边对应图的索引。
batch_graph = tfg.BatchGraph.from_graphs([graph, graph, graph, graph, graph])

# 我们可以将 BatchGraph 对象反向拆分为 Graph 对象
graphs = batch_graph.to_graphs()

# 定义一个图卷积层（GCN）
batch_gcn_layer = tfg.layers.GCN(4, activation=tf.nn.relu)
# 在 BatchGraph 上执行 GCN
batch_h = gcn_layer([batch_graph.x, batch_graph.edge_index, batch_graph.edge_weight])
print("节点表示（BatchGraph 上的 GCN）：\n", batch_h)

# 图池化算法通常依赖于这种批处理数据结构
# 大多数算法接受 BatchGraph 的数据作为输入，并为批次中的每张图输出一个特征向量
graph_h = tfg.nn.mean_pool(batch_h, batch_graph.node_graph_index, num_graphs=batch_graph.num_graphs)
print("图表示（BatchGraph 上的均值池化）：\n", batch_h)


# 定义一个用于对每个节点打分的图卷积层（GCN）
gcn_score_layer = tfg.layers.GCN(1)
# 我们提供了一些高级的图池化操作，例如 topk_pool
node_score = gcn_score_layer([batch_graph.x, batch_graph.edge_index, batch_graph.edge_weight])
node_score = tf.reshape(node_score, [-1])
print("每个节点的得分：\n", node_score)
topk_node_index = tfg.nn.topk_pool(batch_graph.node_graph_index, node_score, ratio=0.6)
print("Top-k 节点索引（Top-k 池化）：\n", topk_node_index)




# ==================================== 内置数据集 ====================================
# 所有图数据都以 numpy 格式存储

# Cora 数据集
graph, (train_index, valid_index, test_index) = tfg.datasets.CoraDataset().load_data()

# PPI 数据集
train_data, valid_data, test_data = tfg.datasets.PPIDataset().load_data()

# TU 数据集
# TU 数据集：https://ls11-www.cs.tu-dortmund.de/staff/morris/graphkerneldatasets
graph_dicts = tfg.datasets.TUDataset("NCI1").load_data()


# ==================================== 基础面向对象 API ====================================
# 面向对象风格的 GCN（图卷积网络）
gcn_layer = tfg.layers.GCN(units=20, activation=tf.nn.relu)

for graph in test_data:
    # 缓存可以通过缓存归一化的边信息来加速 GCN
    outputs = gcn_layer([graph.x, graph.edge_index, graph.edge_weight], cache=graph.cache)
    print(outputs)


# 面向对象风格的 GAT（多头图注意力网络）
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)


# 面向对象风格的多层 GCN 模型
class GCNModel(tf.keras.Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.gcn0 = tfg.layers.GCN(16, activation=tf.nn.relu)
        self.gcn1 = tfg.layers.GCN(7)
        self.dropout = tf.keras.layers.Dropout(0.5)

    def call(self, inputs, training=None, mask=None, cache=None):
        x, edge_index, edge_weight = inputs
        h = self.dropout(x, training=training)
        h = self.gcn0([h, edge_index, edge_weight], cache=cache)
        h = self.dropout(h, training=training)
        h = self.gcn1([h, edge_index, edge_weight], cache=cache)
        return h


gcn_model = GCNModel()
for graph in test_data:
    outputs = gcn_model([graph.x, graph.edge_index, graph.edge_weight], cache=graph.cache)
    print(outputs)


# ==================================== 基础函数式 API ====================================
# 函数式风格的 GCN
# 函数式 API 对于高级算法更加灵活
# 你可以同时传递数据和参数给函数式 API

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, graph.adj(), gcn_w, activation=tf.nn.relu)
    print(outputs)


# ==================================== 高级函数式 API ====================================
# 大多数 API 都是以 Map-Reduce 风格实现的
# 这是一个没有权重归一化和变换的 GCN
# 只需将映射/归约/更新函数传递给函数式 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)

引用

如果您在科学出版物中使用 tf_geometric，我们非常感谢您引用以下论文：

@inproceedings{DBLP:conf/mm/HuQFWZZX21,
  author    = {Jun Hu and
               Shengsheng Qian and
               Quan Fang and
               Youze Wang and
               Quan Zhao and
               Huaiwen Zhang and
               Changsheng Xu},
  editor    = {Heng Tao Shen and
               Yueting Zhuang and
               John R. Smith and
               Yang Yang and
               Pablo Cesar and
               Florian Metze and
               Balakrishnan Prabhakaran},
  title     = {Efficient Graph Deep Learning in TensorFlow with tf{\_}geometric},
  booktitle = {{MM} '21: {ACM} Multimedia Conference, Virtual Event, China, October
               20 - 24, 2021},
  pages     = {3775--3778},
  publisher = {{ACM}},
  year      = {2021},
  url       = {https://doi.org/10.1145/3474085.3478322},
  doi       = {10.1145/3474085.3478322},
  timestamp = {Wed, 20 Oct 2021 12:40:01 +0200},
  biburl    = {https://dblp.org/rec/conf/mm/HuQFWZZX21.bib},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}

相关项目

• MIG-GT： “用于多模态推荐的跨模态全局Transformer图神经网络”（AAAI 2025）。网址：https://github.com/CrawlScript/MIG-GT。
• RpHGNN： “基于随机投影的高效异构图学习”（TKDE 2024）。网址：https://github.com/CrawlScript/RpHGNN。
• MGDCF： “MGDCF：基于马尔可夫图扩散的距离学习用于神经协同过滤”（TKDE 2024）。网址：https://github.com/CrawlScript/Torch-MGDCF。
• tf_sparse： 我们开发了TensorFlow Sparse (tf_sparse)，用于在tf_geometric中实现高效且优雅的稀疏TensorFlow运算。网址：https://github.com/CrawlScript/tf_sparse。
• GRecX： GRecX 是一个高效、统一的基于GNN的推荐基准测试平台。网址：https://github.com/maenzhier/GRecX。
• 🐈 MMClaw： 超轻量级、纯Python的多模态AI智能体核心框架。网址：https://github.com/CrawlScript/MMClaw。

tf_geometric 快速上手指南

tf_geometric 是一个高效且友好的图神经网络（GNN）库，专为 TensorFlow 1.x 和 2.x 设计。它受 pytorch_geometric 启发，采用消息传递机制实现 GNN，相比基于稠密矩阵的实现更高效，比基于稀疏矩阵的实现更易用。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Windows / Linux / Mac OS
• Python 版本: >= 3.7
• 核心依赖:
  • tensorflow 或 tensorflow-gpu: >= 1.15.0 或 >= 2.7.0
  • tf_sparse
  • numpy >= 1.17.4
  • networkx >= 2.1
  • scipy >= 1.1.0

提示：国内用户若遇到下载慢的问题，建议在安装命令前添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像源加速。

安装步骤

根据您的 TensorFlow 版本和需求（CPU/GPU），选择以下任一命令进行安装：

# 仅安装 tf_geometric (需自行预先安装 tensorflow)
pip install -U tf_geometric

# 安装 tf_geometric + TensorFlow 1.x CPU 版
pip install -U tf_geometric[tf1-cpu]

# 安装 tf_geometric + TensorFlow 1.x GPU 版
pip install -U tf_geometric[tf1-gpu]

# 安装 tf_geometric + TensorFlow 2.x CPU 版
pip install -U tf_geometric[tf2-cpu]

# 安装 tf_geometric + TensorFlow 2.x GPU 版
pip install -U tf_geometric[tf2-gpu]

(国内加速示例：pip install -U tf_geometric[tf2-gpu] -i https://pypi.tuna.tsinghua.edu.cn/simple)

基本使用

以下示例展示了如何构建一个图数据对象，并使用多头图注意力网络（GAT）进行处理。这是最基础且典型的使用流程。

# coding=utf-8
import numpy as np
import tf_geometric as tfg
import tensorflow as tf

# 1. 构建图数据
# x: 节点特征 (5 个节点，每个节点 20 维特征)
# edge_index: 边索引 (无向边，格式为 [源节点列表，目标节点列表])
graph = tfg.Graph(
    x=np.random.randn(5, 20),
    edge_index=[[0, 0, 1, 3],
                [1, 2, 2, 1]]
)

print("原始图描述:\n", graph)

# 2. 预处理边
# 将无向边转换为有向边（包含双向），以便 GNN 层处理
graph = graph.to_directed()
print("处理后的边索引:\n", graph.edge_index)

# 3. 定义并运行 GAT 层
# 创建一个 4 头注意力机制的 GAT 层，输出维度为 4，激活函数为 ReLU
gat_layer = tfg.layers.GAT(units=4, num_heads=4, activation=tf.nn.relu)

# 执行前向传播
output = gat_layer([graph.x, graph.edge_index])

print("GAT 输出结果:\n", output)

运行结果示例：

原始图描述:
 Graph Shape: x => (5, 20)	edge_index => (2, 4)	y => None
处理后的边索引:
 [[0 0 1 1 1 2 2 3]
 [1 2 0 2 3 0 1 1]]
GAT 输出结果:
 tf.Tensor(
[[0.22443159 0.         0.58263206 0.32468423]
 [0.29810357 0.         0.19403605 0.35630274]
 ...
 [0.         0.         0.8013462  0.        ]], shape=(5, 4), dtype=float32)

进阶提示

• 批处理: 对于图分类任务，可使用 tfg.BatchGraph.from_graphs() 将多个图打包成一个批次进行处理。
• 缓存加速: 在使用 GCN 等需要计算归一化邻接矩阵的模型时，传入 cache=graph.cache 参数可避免重复计算，显著提升训练速度。
• 内置数据集: 库内置了 Cora、PPI 等常用数据集，可通过 tfg.datasets.CoraDataset().load_data() 直接加载。