in graphlearn_torch/python/partition/base.py [0:0]
def partition(self, with_feature=True, graph_caching=False):
r""" Partition graph and feature data into different parts.
Args:
with_feature (bool): A flag indicating if the feature should be
partitioned with the graph (default: ``True``).
graph_caching (bool): A flag indicating if the full graph topology
will be saved (default: ``False``).
The output directory of partitioned graph data will be like:
* homogeneous
root_dir/
|-- META
|-- node_pb.pt
|-- edge_pb.pt
|-- part0/
|-- graph/
|-- rows.pt
|-- cols.pt
|-- eids.pt
|-- weights.pt (optional)
|-- node_feat/
|-- feats.pkl
|-- ids.pkl
|-- cache_feats.pt (optional)
|-- cache_ids.pt (optional)
|-- edge_feat/
|-- feats.pkl
|-- ids.pkl
|-- cache_feats.pt (optional)
|-- cache_ids.pt (optional)
|-- part1/
|-- graph/
...
|-- node_feat/
...
|-- edge_feat/
...
* heterogeneous
root_dir/
|-- META
|-- node_pb/
|-- ntype1.pt
|-- ntype2.pt
|-- edge_pb/
|-- etype1.pt
|-- etype2.pt
|-- part0/
|-- graph/
|-- etype1/
|-- rows.pt
|-- cols.pt
|-- eids.pt
|-- weights.pt
|-- etype2/
...
|-- node_feat/
|-- ntype1/
|-- feats.pkl
|-- ids.pkl
|-- cache_feats.pt (optional)
|-- cache_ids.pt (optional)
|-- ntype2/
...
|-- edge_feat/
|-- etype1/
|-- feats.pkl
|-- ids.pkl
|-- cache_feats.pt (optional)
|-- cache_ids.pt (optional)
|-- etype2/
...
|-- part1/
|-- graph/
...
|-- node_feat/
...
|-- edge_feat/
...
"""
if 'hetero' == self.data_cls:
node_pb_dict = {}
for ntype in self.node_types:
node_ids_list, node_pb = self._partition_node(ntype)
save_node_pb(self.output_dir, node_pb, ntype)
node_pb_dict[ntype] = node_pb
if with_feature:
self._partition_and_save_node_feat(node_ids_list, ntype)
for etype in self.edge_types:
graph_list, edge_pb = self._partition_graph(node_pb_dict, etype)
edge_feat = self.get_edge_feat(etype)
with_edge_feat = (edge_feat != None)
if graph_caching:
if with_edge_feat:
save_edge_pb(self.output_dir, edge_pb, etype)
save_graph_cache(self.output_dir, graph_list, etype, with_edge_feat)
else:
save_edge_pb(self.output_dir, edge_pb, etype)
for pidx in range(self.num_parts):
save_graph_partition(self.output_dir, pidx, graph_list[pidx], etype)
if with_feature:
self._partition_and_save_edge_feat(graph_list, etype)
else:
node_ids_list, node_pb = self._partition_node()
save_node_pb(self.output_dir, node_pb)
if with_feature:
self._partition_and_save_node_feat(node_ids_list)
graph_list, edge_pb = self._partition_graph(node_pb)
edge_feat = self.get_edge_feat()
with_edge_feat = (edge_feat != None)
if graph_caching:
if with_edge_feat:
save_edge_pb(self.output_dir, edge_pb)
save_graph_cache(self.output_dir, graph_list, with_edge_feat)
else:
save_edge_pb(self.output_dir, edge_pb)
for pidx in range(self.num_parts):
save_graph_partition(self.output_dir, pidx, graph_list[pidx])
if with_feature:
self._partition_and_save_edge_feat(graph_list)
# save meta.
save_meta(self.output_dir, self.num_parts, self.data_cls,
self.node_types, self.edge_types)