Set seed

benchmarks/cugraph-dgl/python-script/dgl_dataloading_benchmark/dgl_benchmark.py

@@ -18,6 +18,8 @@
 import os
 import time
 import json
+import random
+import numpy as np
 from argparse import ArgumentParser
 
 
@@ -218,6 +220,11 @@ def dataloading_benchmark(g, train_idx, fanouts, batch_sizes, use_uva):
  print("==============================================")
  return time_ls
 
+def set_seed(seed):
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
 
 if __name__ == "__main__":
  parser = ArgumentParser()
@@ -230,13 +237,14 @@ def dataloading_benchmark(g, train_idx, fanouts, batch_sizes, use_uva):
  )
  parser.add_argument("--batch_sizes", type=str, default="512,1024")
  parser.add_argument("--do_not_use_uva", action="store_true")
+ parser.add_argument("--seed", type=int, default=42)
  args = parser.parse_args()
 
  if args.do_not_use_uva:
  use_uva = False
  else:
  use_uva = True
-
+ set_seed(args.seed)
  replication_factors = [int(x) for x in args.replication_factors.split(",")]
  fanouts = [[int(y) for y in x.split("_")] for x in args.fanouts.split(",")]
  batch_sizes = [int(x) for x in args.batch_sizes.split(",")]