dmlc--dgl
Mufei Li
7a41c126f6
* Update distributed-preprocessing.rst * Update Co-authored-by: Ubuntu <ubuntu@ip-172-31-9-26.ap-northeast-1.compute.internal>
49 行
1.9 KiB
Python
49 行
1.9 KiB
Python
# Requires setting PYTHONPATH=${GITROOT}/tools
|
|
import json
|
|
import logging
|
|
import sys
|
|
import os
|
|
import numpy as np
|
|
import argparse
|
|
|
|
from utils import setdir
|
|
from utils import array_readwriter
|
|
|
|
def _random_partition(metadata, num_parts):
|
|
num_nodes_per_type = [sum(_) for _ in metadata['num_nodes_per_chunk']]
|
|
ntypes = metadata['node_type']
|
|
for ntype, n in zip(ntypes, num_nodes_per_type):
|
|
logging.info('Generating partition for node type %s' % ntype)
|
|
parts = np.random.randint(0, num_parts, (n,))
|
|
array_readwriter.get_array_parser(name='csv').write(ntype + '.txt', parts)
|
|
|
|
def random_partition(metadata, num_parts, output_path):
|
|
"""
|
|
Randomly partition the graph described in metadata and generate partition ID mapping
|
|
in :attr:`output_path`.
|
|
|
|
A directory will be created at :attr:`output_path` containing the partition ID
|
|
mapping files named "<node-type>.txt" (e.g. "author.txt", "paper.txt" and
|
|
"institution.txt" for OGB-MAG240M). Each file contains one line per node representing
|
|
the partition ID the node belongs to.
|
|
"""
|
|
with setdir(output_path):
|
|
_random_partition(metadata, num_parts)
|
|
|
|
# Run with PYTHONPATH=${GIT_ROOT_DIR}/tools
|
|
# where ${GIT_ROOT_DIR} is the directory to the DGL git repository.
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
'--in_dir', type=str, help='input directory that contains the metadata file')
|
|
parser.add_argument(
|
|
'--out_dir', type=str, help='output directory')
|
|
parser.add_argument(
|
|
'--num_partitions', type=int, help='number of partitions')
|
|
logging.basicConfig(level='INFO')
|
|
args = parser.parse_args()
|
|
with open(os.path.join(args.in_dir, 'metadata.json')) as f:
|
|
metadata = json.load(f)
|
|
num_parts = args.num_partitions
|
|
random_partition(metadata, num_parts, args.out_dir)
|