GraphStore & FeatureStore 技术分享
1.引言
图结构数据由两种类型的数据组成：
图拓扑：图中的节点以及它们之间的连边 节点和边的特征：对应于图中节点和边的属性的密集向量
GNN使用图结构数据进行训练，当图数据过大时，通常需要用采样的方式获得mini-batch，在mini-batch上进行训练。但是mini-batch训练仍然需要将图拓扑和特征放在内存上，无法扩展到规模比内存大的图。因此，想要将GNN扩展到超出单台机器内存要求的大规模图，需要将图拓扑和特征移出内存。
2.GraphStore & FeatureStore
将图拓扑和特征移出内存的一种方案是将图数据存入数据库，通过与数据库交互提取图拓扑和相应的特征。GammaGL通过 GraphStore 和 FeatureStore 接口来实现这个功能，GraphStore 存取图拓扑，FeatureStore 存取节点特征。
通常情况下图数据中特征的数据量往往比图拓扑大好几个数量级，所以特征存储是大图训练时的主要存储瓶颈。FeatureStore 抽象出了对图特征进行增删改查的操作，可以根据使用的数据库，主要重写_put_tensor() 、 _get_tensor() 和 _remove_tensor() 方法实现相应数据库对节点特征进行增删改查的操作。同时，FeatureStore提供了简便直观的接口操作FeatureStore中的数据，如下：
feature_store = MyFeatureStore()
paper_features = ... # [num_papers, num_paper_features]
# Add features
feature_store['paper', 'attribute', [1,2,3]] = paper_features
# Access features
assert torch.equal(feature_store['paper', 'attribute'], paper_features)
# Del features
del feature_store['paper', 'attribute', [1,2,3]]
GraphStore 抽象出了对图拓扑进行增删改查的操作，可以根据使用的数据库，主要重写 _put_edge_index() 、_get_edge_index() 和 _remove_edge_index() 方法实现相应数据库对边索引进行增删改查的操作。GraphStore 支持 coo、csc、csr 格式的边，并提供进行类型转换的接口。同样的，GraphStore也可以通过简便的接口操作，如下：
graph_store = MyGraphStore()
edge = ... # [2, num_edges]
# Add edges
graph_store['cite', 'coo'] = edge
# Access edges
assert torch.equal(graph_store['cite', 'coo'], edge)
# Del edges
del graph_store['cite', 'coo']
3.数据库实例
GammaGL提供了使用 graphstore 和 featurestore 存储reddit数据集并将数据库中的数据用于训练的实例。
我们基于 neo4j 和 NebulaGraph 两个数据库分别实现了相应的 featureftore 和 graphstore，训练时先从 graphstore 中取出图拓扑，利用采样器对大图采样得到子图，然后再根据子图的 id 从 featurestore 中取出对应的节点特征。这样就可以避免把用不到的节点特征取入内存，极大的减少了内存占用。
Neo4jGraphStore 核心代码如下：
class Neo4jGraphStore(GraphStore):
def __init__(self, uri: str, user_name: str, password: str, **kwargs):
super().__init__()
self.driver = GraphDatabase.driver(uri, auth=(user_name, password))
def __del__(self):
self.driver.close()
def _put_edge_index(self, edge_index, edge_attr: EdgeAttr) -> bool:
batch_size = 128
with self.driver.session() as session:
for i in range(0, edge_index[0].size(0), batch_size):
batch_queries = []
for j in range(i, min(i + batch_size, edge_index[0].size(0))):
batch_queries.append({
"src_id": edge_index[0][j].item(),
"dst_id": edge_index[1][j].item()
})
session.run(
"""
UNWIND $rels AS rel
MERGE (n {ID: rel.src_id})
MERGE (m {ID: rel.dst_id})
""" + f"MERGE (n)-[r:{edge_attr.edge_type}]->(m)",
rels=batch_queries
)
return True
def _get_edge_index(self, edge_attr: EdgeAttr) -> EdgeTensorType:
with self.driver.session() as session:
result = session.run(
f"MATCH (n)-[r:{edge_attr.edge_type}]->(m) RETURN n.ID AS src_id, m.ID AS dst_id",
)
src_ids = []
dst_ids = []
for record in result:
src_ids.append(int(record['src_id']))
dst_ids.append(int(record['dst_id']))
edge_index = tlx.convert_to_tensor([src_ids,dst_ids])
return edge_index
def _remove_edge_index(self, edge_attr: EdgeAttr) -> bool:
with self.driver.session() as session:
session.run(f"MATCH (n)-[r:{edge_attr.edge_type}]->(m) DELETE r")
return True
def get_all_edge_attrs(self) -> List[EdgeAttr]:
with self.driver.session() as session:
result = session.run("CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType")
return [EdgeAttr(record['relationshipType'], EdgeLayout.COO) for record in result]
更多的实现可以在 https://github.com/xy-Ji/gdbi 查看。
我们的实现只是示例，用户可以在更多数据库上实现 featurestore 和 graphstore，也可以添加更多的功能或提供更高效的实现。如在键值数据库上实现featurestore，我们的示例实现是在图数据库上同时实现了graphstore和featurestore，事实上完全可以将它们分别在图数据库和键值数据库上实现，键值数据库根据 id 查询属性的效率远大于图数据库，因此基于键值数据库实现featurestore可以极大的提升性能。
实验室算法库简介
两算法库均获得OpenI启智社区“社区优秀孵化项目奖”
GammaGL：
GammaGL是一个基于TensorLayerX实现的开源的多框架的图神经网络（GNN）算法平台，该框架抽象出图数据、消息传递、采样等组件，已经实现了二十余种图数据集和六十余种术界经典GNN算法。GammaGL发布0.5版本 https://github.com/BUPT-GAMMA/GammaGL
OpenHGNN：
OpenHGNN是一个基于深度学习框架PyTorch和图神经网络框架DGL作为底层框架的异质图神经网络（HGNN）工具包。该工具包专注于HGNN的算法模型，抽象出了数据集管理、算法流程、任务评测等模块，提供了超参数调优组件和最优的超参数设置，发布了HGNN设计空间和Benchmark等。OpenHGNN发布0.7版本OpenHGNN：https://github.com/BUPT-GAMMA/OpenHGNN
