PyTorch Cluster

To build from source, ensure PyTorch 1.4.0 or newer is installed and that cuda/bin and cuda/include are in your $PATH and $CPATH respectively.

Docker/Non-NVIDIA Driver Note: If running in a container without an NVIDIA driver, PyTorch may fail to evaluate compute capabilities. Set the compute capabilities manually using the TORCH_CUDA_ARCH_LIST environment variable.

export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.2+PTX 7.5+PTX"
pip install torch-cluster

Pre-built pip wheels are available for various OS, PyTorch, and CUDA combinations. To install, use the -f flag pointing to the appropriate URL for your PyTorch version.

For PyTorch 2.11:

pip install torch-cluster -f https://data.pyg.org/whl/torch-2.11.0+${CUDA}.html

For PyTorch 2.10:

pip install torch-cluster -f https://data.pyg.org/whl/torch-2.10.0+${CUDA}.html

Replace ${CUDA} with one of the following based on your installation:

• cpu
• cu126
• cu128
• cu130

Note: Binaries are also available for older PyTorch versions (1.4.0 through 2.8.0). Check the official wheel index for the latest supported version number.

The torch-cluster package offers a C++ API that provides C++ equivalents of the Python models. To build it, use CMake and ensure Torch_DIR is set correctly.

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

The graclus_cluster function implements a greedy clustering algorithm that picks an unmarked vertex and matches it with an unmarked neighbor that maximizes edge weight.

Arguments:

• row (Tensor): Row indices of edges.
• col (Tensor): Column indices of edges.
• weight (Tensor, optional): Edge weights.

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

cluster = graclus_cluster(row, col, weight)

Computes graph edges to the nearest k points.

Arguments:

• x (Tensor): Node feature matrix of shape [N, F].
• k (int): Number of neighbors.
• batch (LongTensor, optional): Batch vector of shape [N]. Must be sorted.
• loop (bool, optional): If True, includes self-loops. (default: False)
• flow (string, optional): Flow direction ("source_to_target" or "target_to_source"). (default: "source_to_target")
• cosine (bool, optional): If True, uses Cosine distance instead of Euclidean. (default: False)
• num_workers (int, optional): Number of workers for computation. (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)

Samples random walks of length walk_length from all node indices in start in the graph defined by (row, col).

Arguments:

• row (Tensor): Edge row indices.
• col (Tensor): Edge column indices.
• start (Tensor): Starting node indices.
• walk_length (int): Length of the walks.

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)

The fps function iteratively samples the most distant point with regard to the rest points.

Arguments:

• x (Tensor): Input points.
• batch (LongTensor): Batch vector assigning nodes to examples (must be sorted).
• ratio (float): Sampling ratio.
• random_start (bool): Whether to start from a random point.

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)

Computes graph edges to all points within a given distance r.

Arguments:

• x (Tensor): Node feature matrix of shape [N, F].
• r (float): The radius.
• batch (LongTensor, optional): Batch vector of shape [N]. Must be sorted.
• loop (bool, optional): If True, includes self-loops. (default: False)
• max_num_neighbors (int, optional): Maximum neighbors to return per element. If exceeded, neighbors are picked randomly. (default: 32)
• flow (string, optional): Flow direction ("source_to_target" or "target_to_source"). (default: "source_to_target")
• num_workers (int, optional): Number of workers for computation. (default: 1)

import torch
from torch_cluster import radius_graph

x = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]])
edge_index = radius_graph(x, r=2.5, batch=None, loop=False)

Clusters points in x together which are nearest to a given query point in y.

Note: batch_x and batch_y vectors must 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)

The grid_cluster function overlays a regular grid of a user-defined size over a point cloud and clusters all points within a voxel.

Arguments:

• pos (Tensor): Point cloud positions.
• size (Tensor): The size of the grid voxels.

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)