Benchmarks for biased sampling

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [0.4.0] - 2023-MM-DD
 ### Added
 ### Changed
+- Added `--biased` parameter to run benchmarks for biased sampling ([#267](https://github.com/pyg-team/pyg-lib/pull/267))
 ### Removed
 
 ## [0.3.0] - 2023-10-11

benchmark/sampler/hetero_neighbor.py

@@ -33,6 +33,7 @@
 # TODO(kgajdamo): Enable sampling with replacement
 # argparser.add_argument('--replace', action='store_true')
 argparser.add_argument('--shuffle', action='store_true')
+argparser.add_argument('--biased', action='store_true')
 argparser.add_argument('--temporal', action='store_true')
 argparser.add_argument('--temporal-strategy', choices=['uniform', 'last'],
                        default='uniform')
@@ -51,6 +52,7 @@ def test_hetero_neighbor(dataset, **kwargs):
 
     colptr_dict, row_dict = dataset
     num_nodes_dict = {k[-1]: v.size(0) - 1 for k, v in colptr_dict.items()}
+    num_edges_dict = {k: v.size(0) for k, v in row_dict.items()}
 
     if args.temporal:
         # generate random timestamps
@@ -60,6 +62,13 @@ def test_hetero_neighbor(dataset, **kwargs):
     else:
         node_time_dict = None
 
+    edge_weight_dict = None
+    if args.biased:
+        edge_weight_dict = {
+            edge_type: torch.rand(num_edges)
+            for edge_type, num_edges in num_edges_dict.items()
+        }
+
     if args.shuffle:
         node_perm = torch.randperm(num_nodes_dict['paper'])
     else:
@@ -86,6 +95,7 @@ def test_hetero_neighbor(dataset, **kwargs):
                     num_neighbors_dict,
                     node_time_dict,
                     seed_time_dict=None,
+                    edge_weight_dict=edge_weight_dict,
                     csc=True,
                     replace=False,
                     directed=True,

benchmark/sampler/neighbor.py

@@ -29,6 +29,7 @@
 ])
 argparser.add_argument('--replace', action='store_true')
 argparser.add_argument('--shuffle', action='store_true')
+argparser.add_argument('--biased', action='store_true')
 argparser.add_argument('--temporal', action='store_true')
 argparser.add_argument('--temporal-strategy', choices=['uniform', 'last'],
                        default='uniform')
@@ -45,7 +46,10 @@ def test_neighbor(dataset, **kwargs):
         raise ValueError(
             "Temporal sampling needs to create disjoint subgraphs")
 
-    (rowptr, col), num_nodes = dataset, dataset[0].size(0) - 1
+    rowptr, col = dataset
+    num_nodes = rowptr.numel() - 1
+    num_edges = col.numel()
+
     if 'dgl' in args.libraries:
         import dgl
         dgl_graph = dgl.graph(
@@ -57,6 +61,10 @@ def test_neighbor(dataset, **kwargs):
     else:
         node_time = None
 
+    edge_weight = None
+    if args.biased:
+        edge_weight = torch.rand(num_edges)
+
     if args.shuffle:
         node_perm = torch.randperm(num_nodes)
     else:
@@ -80,6 +88,7 @@ def test_neighbor(dataset, **kwargs):
                     num_neighbors,
                     time=node_time,
                     seed_time=None,
+                    edge_weight=edge_weight,
                     replace=args.replace,
                     directed=args.directed,
                     disjoint=args.disjoint,