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:
All included operations work on varying data types and are implemented both for CPU and GPU.
Installation AnacondaUpdate: You can now install pytorch-cluster
via Anaconda for all major OS/PyTorch/CUDA combinations ๐ค Given that you have pytorch >= 1.8.0
installed, simply run
conda install pytorch-cluster -c pyg
Binaries
We alternatively provide pip wheels for all major OS/PyTorch/CUDA combinations, see here.
PyTorch 2.1To install the binaries for PyTorch 2.1.0, simply run
pip install torch-cluster -f https://data.pyg.org/whl/torch-2.1.0+${CUDA}.html
where ${CUDA}
should be replaced by either cpu
, cu118
, or cu121
depending on your PyTorch installation.
cpu
cu118
cu121
Linux โ
โ
โ
Windows โ
โ
โ
macOS โ
PyTorch 2.0
To install the binaries for PyTorch 2.0.0, simply run
pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.html
where ${CUDA}
should be replaced by either cpu
, cu117
, or cu118
depending on your PyTorch installation.
cpu
cu117
cu118
Linux โ
โ
โ
Windows โ
โ
โ
macOS โ
Note: Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1 and PyTorch 1.13.0/1.13.1 (following the same procedure). For older versions, you need to explicitly specify the latest supported version number or install via pip install --no-index
in order to prevent a manual installation from source. You can look up the latest supported version number here.
Ensure that at least PyTorch 1.4.0 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__)"
>>> 1.4.0
$ python -c "import torch; print(torch.__version__)"
>>> 1.1.0
$ echo $PATH
>>> /usr/local/cuda/bin:...
$ echo $CPATH
>>> /usr/local/cuda/include:...
Then run:
pip install torch-cluster
When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail. In this case, ensure that the compute capabilities are set via TORCH_CUDA_ARCH_LIST
, e.g.:
export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"
Functions Graclus
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])
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])
FarthestPointSampling
A sampling algorithm, which iteratively samples the most distant point with regard to the rest points.
import torch from torch_cluster import fps x = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]]) batch = torch.tensor([0, 0, 0, 0]) index = fps(x, batch, ratio=0.5, random_start=False)
print(index)
tensor([0, 3])
kNN-Graph
Computes graph edges to the nearest k points.
Args:
[N, F]
.[N]
, which assigns each node to a specific example. batch
needs to be sorted. (default: None
)True
, the graph will contain self-loops. (default: False
)"source_to_target"
or "target_to_source"
). (default: "source_to_target"
)True
, will use the Cosine distance instead of Euclidean distance to find nearest neighbors. (default: False
)batch
is not None
, or the input lies on the GPU. (default: 1
)import torch from torch_cluster import knn_graph x = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]]) batch = torch.tensor([0, 0, 0, 0]) edge_index = knn_graph(x, k=2, batch=batch, loop=False)
print(edge_index)
tensor([[1, 2, 0, 3, 0, 3, 1, 2],
[0, 0, 1, 1, 2, 2, 3, 3]])
Radius-Graph
Computes graph edges to all points within a given distance.
Args:
[N, F]
.[N]
, which assigns each node to a specific example. batch
needs to be sorted. (default: None
)True
, the graph will contain self-loops. (default: False
)max_num_neighbors
, returned neighbors are picked randomly. (default: 32
)"source_to_target"
or "target_to_source"
). (default: "source_to_target"
)batch
is not None
, or the input lies on the GPU. (default: 1
)import torch from torch_cluster import radius_graph x = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]]) batch = torch.tensor([0, 0, 0, 0]) edge_index = radius_graph(x, r=2.5, batch=batch, loop=False)
print(edge_index)
tensor([[1, 2, 0, 3, 0, 3, 1, 2],
[0, 0, 1, 1, 2, 2, 3, 3]])
Nearest
Clusters points in x together which are nearest to a given query point in y. batch_{x,y}
vectors need to be sorted.
import torch from torch_cluster import nearest x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) batch_x = torch.tensor([0, 0, 0, 0]) y = torch.Tensor([[-1, 0], [1, 0]]) batch_y = torch.tensor([0, 0]) cluster = nearest(x, y, batch_x, batch_y)
print(cluster)
tensor([0, 0, 1, 1])
RandomWalk-Sampling
Samples random walks of length walk_length
from all node indices in start
in the graph given by (row, col)
.
import torch from torch_cluster import random_walk row = torch.tensor([0, 1, 1, 1, 2, 2, 3, 3, 4, 4]) col = torch.tensor([1, 0, 2, 3, 1, 4, 1, 4, 2, 3]) start = torch.tensor([0, 1, 2, 3, 4]) walk = random_walk(row, col, start, walk_length=3)
print(walk)
tensor([[0, 1, 2, 4],
[1, 3, 4, 2],
[2, 4, 2, 1],
[3, 4, 2, 4],
[4, 3, 1, 0]])
Running tests
pytest
C++ API
torch-cluster
also offers a C++ API that contains C++ equivalent of python models.
export Torch_DIR=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`
mkdir build
cd build
# Add -DWITH_CUDA=on support for the CUDA if needed
cmake ..
make
make install
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4