This package consists of a small extension library of highly optimized graph cluster algorithms for the use in PyTorch. The package consists of the following clustering algorithms:
- Graclus from Dhillon et al.: Weighted Graph Cuts without Eigenvectors: A Multilevel Approach (PAMI 2007)
- Voxel Grid Pooling from, e.g., Simonovsky and Komodakis: Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs (CVPR 2017)
All included operations work on varying data types and are implemented both for CPU and GPU.
Ensure that at least PyTorch 0.4.1 is installed and verify that cuda/bin and cuda/include are in your $PATH and $CPATH respectively, e.g.:
$ python -c "import torch; print(torch.__version__)"
>>> 0.4.1
$ echo $PATH
>>> /usr/local/cuda/bin:...
$ echo $CPATH
>>> /usr/local/cuda/include:...
Then run:
pip install torch-cluster
If you are running into any installation problems, please create an issue.
Be sure to import torch first before using this package to resolve symbols the dynamic linker must see.
A greedy clustering algorithm of picking an unmarked vertex and matching it with one its unmarked neighbors (that maximizes its edge weight). The GPU algorithm is adapted from Fagginger Auer and Bisseling: A GPU Algorithm for Greedy Graph Matching (LNCS 2012)
import torch
from torch_cluster import graclus_cluster
row = torch.tensor([0, 1, 1, 2])
col = torch.tensor([1, 0, 2, 1])
weight = torch.Tensor([1, 1, 1, 1]) # Optional edge weights.
cluster = graclus_cluster(row, col, weight)print(cluster)
tensor([0, 0, 1])
A clustering algorithm, which overlays a regular grid of user-defined size over a point cloud and clusters all points within a voxel.
import torch
from torch_cluster import grid_cluster
pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
size = torch.Tensor([5, 5])
cluster = grid_cluster(pos, size)print(cluster)
tensor([0, 5, 3, 0, 1])
python setup.py test