add CHANGELOG; fix unit test error; add train_idx and test_idx partit…

…ion in example

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

CHANGELOG.md

@@ -51,6 +51,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for `torch.jit.script` within `MessagePassing` layers without `torch_sparse` being installed ([#7061](https://github.com/pyg-team/pytorch_geometric/pull/7061), [#7062](https://github.com/pyg-team/pytorch_geometric/pull/7062))
 - Added unbatching logic for `torch.sparse` tensors ([#7037](https://github.com/pyg-team/pytorch_geometric/pull/7037))
 - Added the `RotatE` KGE model ([#7026](https://github.com/pyg-team/pytorch_geometric/pull/7026))
+- Add partitioning for distributed training ([#7502](https://github.com/pyg-team/pytorch_geometric/pull/7502))
 
 ### Changed
 

examples/distributed/partition_graph.py

@@ -1,9 +1,13 @@
 import argparse
 import os.path as osp
 
+import torch
 from ogb.nodeproppred import PygNodePropPredDataset
 
-from torch_geometric.distributed.partition.partitioner import Partitioner
+from torch_geometric.distributed.partition.partitioner import (
+    Partitioner,
+    prepare_directory,
+)
 
 
 def partition_dataset(ogbn_dataset: str, root_dir: str, num_partitions: int):
@@ -14,6 +18,28 @@ def partition_dataset(ogbn_dataset: str, root_dir: str, num_partitions: int):
     partitioner = Partitioner(output_dir=save_dir, num_parts=num_partitions,
                               data=data)
     partitioner.generate_partition()
+    split_idx = dataset.get_idx_split()
+    print('-- Saving label ...')
+    label_dir = prepare_directory(root_dir, f'{ogbn_dataset}-label')
+    torch.save(data.y.squeeze(), osp.join(label_dir, 'label.pt'))
+
+    print('-- Partitioning training idx ...')
+    train_idx = split_idx['train']
+    train_idx = train_idx.split(train_idx.size(0) // num_partitions)
+    train_idx_partitions_dir = prepare_directory(
+        root_dir, f'{ogbn_dataset}-train-partitions')
+    for pidx in range(num_partitions):
+        torch.save(train_idx[pidx],
+                   osp.join(train_idx_partitions_dir, f'partition{pidx}.pt'))
+
+    print('-- Partitioning test idx ...')
+    test_idx = split_idx['test']
+    test_idx = test_idx.split(test_idx.size(0) // num_partitions)
+    test_idx_partitions_dir = prepare_directory(
+        root_dir, f'{ogbn_dataset}-test-partitions')
+    for pidx in range(num_partitions):
+        torch.save(test_idx[pidx],
+                   osp.join(test_idx_partitions_dir, f'partition{pidx}.pt'))
 
 
 if __name__ == '__main__':

test/distributed/test_partition.py

@@ -4,9 +4,11 @@
 
 from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
 from torch_geometric.distributed.partition.partitioner import Partitioner
+from torch_geometric.testing import withPackage
 from torch_geometric.typing import EdgeTypeStr
 
 
+@withPackage('pyg_lib')
 def test_partition_homo_graph():
     num_partitions = 2
     save_dir = osp.join('./homo', 'partition')
@@ -58,6 +60,7 @@ def test_partition_homo_graph():
     assert torch.allclose(data.x[node_id2], node_feat2)
 
 
+@withPackage('pyg_lib')
 def test_partition_hetero_graph():
     num_partitions = 2
     save_dir = osp.join('./hetero', 'partition')