so much faster

benchmark/benchmark.py

@@ -0,0 +1,31 @@
+import time
+
+import torch
+from torch_cluster import sparse_grid_cluster
+
+n = 90000000
+s = 1 / 64
+
+print('GPU ===================')
+
+t = time.perf_counter()
+pos = torch.cuda.FloatTensor(n, 3).uniform_(0, 1)
+size = torch.cuda.FloatTensor([s, s, s])
+torch.cuda.synchronize()
+print('Init:', time.perf_counter() - t)
+
+t_all = time.perf_counter()
+sparse_grid_cluster(pos, size)
+torch.cuda.synchronize()
+t_all = time.perf_counter() - t_all
+print('All:', t_all)
+
+print('CPU ===================')
+
+pos = pos.cpu()
+size = size.cpu()
+
+t_all = time.perf_counter()
+sparse_grid_cluster(pos, size)
+t_all = time.perf_counter() - t_all
+print('All:', t_all)

setup.py

@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '0.2.3'
+__version__ = '0.2.4'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['cffi', 'torch-unique']

torch_cluster/__init__.py

@@ -1,5 +1,5 @@
 from .functions.grid import sparse_grid_cluster, dense_grid_cluster
 
-__version__ = '0.2.3'
+__version__ = '0.2.4'
 
 __all__ = ['sparse_grid_cluster', 'dense_grid_cluster', '__version__']

torch_cluster/functions/grid.py

@@ -18,11 +18,12 @@ def _preprocess(position, size, batch=None, start=None):
 
     # Translate to minimal positive positions if no start was passed.
     if start is None:
-        position = position - position.min(dim=-2, keepdim=True)[0]
-    else:
+        min = []
+        for i in range(position.size(-1)):
+            min.append(position[:, i].min())
+        position = position - position.new(min)
+    elif start != 0:
         position = position - start
-        assert position.min() >= 0, (
-            'Passed origin resulting in unallowed negative positions')
 
     # If given, append batch to position tensor.
     if batch is not None:
@@ -37,10 +38,10 @@ def _preprocess(position, size, batch=None, start=None):
 
 
 def _minimal_cluster_size(position, size):
-    max = position.max(dim=0)[0]
-    while max.dim() > 1:
-        max = max.max(dim=0)[0]
-    cluster_size = (max / size).long() + 1
+    max = []
+    for i in range(position.size(-1)):
+        max.append(position[:, i].max())
+    cluster_size = (size.new(max) / size).long() + 1
     return cluster_size