PyTorch Cluster

---

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
• VoxelGrid

All included operations work on varying data types and are implemented both for CPU and GPU.

Installation

Check that nvcc is accessible from terminal, e.g. nvcc --version. If not, add cuda (/usr/local/cuda/bin) to your $PATH. Then run:

pip install cffi torch-cluster

Graclus

A greedy clustering algorithm of picking an unmarked vertex and matching it with one its unmarked neighbors (that maximizes its edge weight).

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])

VoxelGrid

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])

Running tests

python setup.py test