added performance script and small fixes to readme

Author: Thomas Klijnsma

README.md

@@ -94,7 +94,7 @@ Example Singularity instructions:
 
 ```bash
 singularity pull docker://pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
-singularity run --nv sifs/pytorch_2.0.0-cuda11.7-cudnn8-devel.sif
+singularity run --nv pytorch_2.0.0-cuda11.7-cudnn8-devel.sif
 ```
 
 And then once in the container:

scripts/performance.py

@@ -0,0 +1,41 @@
+import time
+import torch
+import torch_cmspepr
+import torch_cluster
+import tqdm
+gpu = torch.device('cuda')
+
+k = 10
+
+def gen(cuda=False):
+    # 10k nodes with 5 node features
+    x = torch.rand((10000, 5))
+    # Split nodes over 4 events with 2500 nodes/evt
+    batch = torch.repeat_interleave(torch.arange(4), 2500)
+    if cuda: x, batch = x.to(gpu), batch.to(gpu)
+    return x, batch
+
+def profile(name, unit):
+    t0 = time.time()
+    for _ in tqdm.tqdm(range(10)): unit()
+    print(f'{name} took {(time.time() - t0)/100.} sec/evt')
+
+def cpu_cmspepr():
+    x, batch = gen()
+    torch_cmspepr.knn_graph(x, k, batch=batch)
+profile('CPU (torch_cmspepr)', cpu_cmspepr)
+
+def cpu_cluster():
+    x, batch = gen()
+    torch_cluster.knn_graph(x, k, batch=batch)
+profile('CPU (torch_cluster)', cpu_cmspepr)
+
+def cuda_cmspepr():
+    x, batch = gen(cuda=True)
+    torch_cmspepr.knn_graph(x, k, batch=batch)
+profile('CUDA (torch_cmspepr)', cuda_cmspepr)
+
+def cuda_cluster():
+    x, batch = gen(cuda=True)
+    torch_cluster.knn_graph(x, k, batch=batch)
+profile('CUDA (torch_cluster)', cpu_cmspepr)