Package: torch-cluster
PyTorch ClusterThis package consists of a small extension library of highly optimized graph cluster algorithms for the use inPyTorch.
The package consists of the following clustering algorithms:Graclusfrom Dhillonet al.:Weighted Graph Cuts without Eigenvectors: A Multilevel Approach(PAMI 2007)Voxel Grid Poolingfrom,e.g., Simonovsky and Komodakis:Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs(CVPR 2017)Iterative Farthest Point Samplingfrom,e.g.Qiet al.:PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space(NIPS 2017)k-NNandRadiusgraph generationClustering based onNearestpointsRandom Walk Samplingfrom,e.g., Grover and Leskovec:node2vec: Scalable Feature Learning for Networks(KDD 2016)All included operations work on varying data types and are implemented both for CPU and GPU.InstallationAnacondaUpdate:You can now installpytorch-clusterviaAnacondafor all major OS/PyTorch/CUDA combinations 🤗
Given that you havepytorch >= 1.8.0installed, simply runconda install pytorch-cluster -c pygBinariesWe alternatively provide pip wheels for all major OS/PyTorch/CUDA combinations, seehere.PyTorch 2.1To install the binaries for PyTorch 2.1.0, simply runpip install torch-cluster -f https://data.pyg.org/whl/torch-2.1.0+${CUDA}.htmlwhere${CUDA}should be replaced by eithercpu,cu118, orcu121depending on your PyTorch installation.cpucu118cu121Linux✅✅✅Windows✅✅✅macOS✅PyTorch 2.0To install the binaries for PyTorch 2.0.0, simply runpip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+${CUDA}.htmlwhere${CUDA}should be replaced by eithercpu,cu117, orcu118depending on your PyTorch installation.cpucu117cu118Linux✅✅✅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 viapip install --no-indexin order to prevent a manual installation from source.
You can look up the latest supported version numberhere.From sourceEnsure that at least PyTorch 1.4.0 is installed and verify thatcuda/binandcuda/includeare in your$PATHand$CPATHrespectively,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-clusterWhen 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 viaTORCH_CUDA_ARCH_LIST,e.g.:export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"FunctionsGraclusA 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)importtorchfromtorch_clusterimportgraclus_clusterrow=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])VoxelGridA clustering algorithm, which overlays a regular grid of user-defined size over a point cloud and clusters all points within a voxel.importtorchfromtorch_clusterimportgrid_clusterpos=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])FarthestPointSamplingA sampling algorithm, which iteratively samples the most distant point with regard to the rest points.importtorchfromtorch_clusterimportfpsx=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-GraphComputes graph edges to the nearestkpoints.Args:x(Tensor): Node feature matrix of shape[N, F].k(int): The number of neighbors.batch(LongTensor, optional): Batch vector of shape[N], which assigns each node to a specific example.batchneeds to be sorted. (default:None)loop(bool, optional): IfTrue, the graph will contain self-loops. (default:False)flow(string, optional): The flow direction when using in combination with message passing ("source_to_target"or"target_to_source"). (default:"source_to_target")cosine(boolean, optional): IfTrue, will use the Cosine distance instead of Euclidean distance to find nearest neighbors. (default:False)num_workers(int): Number of workers to use for computation. Has no effect in casebatchis notNone, or the input lies on the GPU. (default:1)importtorchfromtorch_clusterimportknn_graphx=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-GraphComputes graph edges to all points within a given distance.Args:x(Tensor): Node feature matrix of shape[N, F].r(float): The radius.batch(LongTensor, optional): Batch vector of shape[N], which assigns each node to a specific example.batchneeds to be sorted. (default:None)loop(bool, optional): IfTrue, the graph will contain self-loops. (default:False)max_num_neighbors(int, optional): The maximum number of neighbors to return for each element. If the number of actual neighbors is greater thanmax_num_neighbors, returned neighbors are picked randomly. (default:32)flow(string, optional): The flow direction when using in combination with message passing ("source_to_target"or"target_to_source"). (default:"source_to_target")num_workers(int): Number of workers to use for computation. Has no effect in casebatchis notNone, or the input lies on the GPU. (default:1)importtorchfromtorch_clusterimportradius_graphx=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]])NearestClusters points inxtogether which are nearest to a given query point iny.batch_{x,y}vectors need to be sorted.importtorchfromtorch_clusterimportnearestx=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-SamplingSamples random walks of lengthwalk_lengthfrom all node indices instartin the graph given by(row, col).importtorchfromtorch_clusterimportrandom_walkrow=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 testspytestC++ APItorch-clusteralso 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