README.md 6.37 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
6
7
[pypi-image]: https://badge.fury.io/py/torch-cluster.svg
[pypi-url]: https://pypi.python.org/pypi/torch-cluster
[build-image]: https://travis-ci.org/rusty1s/pytorch_cluster.svg?branch=master
[build-url]: https://travis-ci.org/rusty1s/pytorch_cluster
[coverage-image]: https://codecov.io/gh/rusty1s/pytorch_cluster/branch/master/graph/badge.svg
[coverage-url]: https://codecov.io/github/rusty1s/pytorch_cluster?branch=master

rusty1s's avatar
rusty1s committed
8
# PyTorch Cluster
rusty1s's avatar
rusty1s committed
9

rusty1s's avatar
rusty1s committed
10
11
12
13
[![PyPI Version][pypi-image]][pypi-url]
[![Build Status][build-image]][build-url]
[![Code Coverage][coverage-image]][coverage-url]

rusty1s's avatar
rusty1s committed
14
15
--------------------------------------------------------------------------------

rusty1s's avatar
rusty1s committed
16
This package consists of a small extension library of highly optimized graph cluster algorithms for the use in [PyTorch](http://pytorch.org/).
rusty1s's avatar
typo  
rusty1s committed
17
The package consists of the following clustering algorithms:
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
credit  
rusty1s committed
19
* **[Graclus](#graclus)** from Dhillon *et al.*: [Weighted Graph Cuts without Eigenvectors: A Multilevel Approach](http://www.cs.utexas.edu/users/inderjit/public_papers/multilevel_pami.pdf) (PAMI 2007)
rusty1s's avatar
rusty1s committed
20
* **[Voxel Grid Pooling](#voxelgrid)** from, *e.g.*, Simonovsky and Komodakis: [Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://arxiv.org/abs/1704.02901) (CVPR 2017)
rusty1s's avatar
rusty1s committed
21
22
* **[Iterative Farthest Point Sampling](#farthestpointsampling)** from, *e.g.* Qi *et al.*: [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://arxiv.org/abs/1706.02413) (NIPS 2017)
* **[k-NN](#knn-graph)** and **[Radius](#radius-graph)** graph generation
rusty1s's avatar
rusty1s committed
23
* Clustering based on **[Nearest](#nearest)** points
rusty1s's avatar
rusty1s committed
24
* **[Random Walk Sampling](#randomwalk-sampling)** from, *e.g.*, Grover and Leskovec: [node2vec: Scalable Feature Learning for Networks](https://arxiv.org/abs/1607.00653) (KDD 2016)
rusty1s's avatar
rusty1s committed
25
26
27
28
29

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

## Installation

rusty1s's avatar
rusty1s committed
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
### Binaries

We provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://s3.eu-central-1.amazonaws.com/pytorch-geometric.com/whl/index.html).
To install from binaries, simply run

```
pip install torch-cluster==latest+${CUDA} -f https://s3.eu-central-1.amazonaws.com/pytorch-geometric.com/whl/torch-1.4.0.html
```

where `${CUDA}` should be replaced by either `cpu`, `cu92`, `cu100` or `cu101` depending on your PyTorch installation.

|             | `cpu` | `cu92` | `cu100` | `cu101` |
|-------------|-------|--------|---------|---------|
| **Linux**   | ✅    | ✅     | ✅      | ✅      |
| **Windows** | ✅    | ❌     | ❌      | ✅      |
| **macOS**   | ✅    |        |         |         |

### From source

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
rusty1s's avatar
rusty1s committed
54
55
56

```
$ python -c "import torch; print(torch.__version__)"
rusty1s's avatar
rusty1s committed
57
>>> 1.1.0
rusty1s's avatar
rusty1s committed
58
59
60
61
62

$ echo $PATH
>>> /usr/local/cuda/bin:...

$ echo $CPATH
rusty1s's avatar
rusty1s committed
63
>>> /usr/local/cuda/include:...
rusty1s's avatar
rusty1s committed
64
65
```

rusty1s's avatar
rusty1s committed
66
67
68
Then run:

```
rusty1s's avatar
rusty1s committed
69
pip install torch-cluster
rusty1s's avatar
rusty1s committed
70
71
```

rusty1s's avatar
rusty1s committed
72
73
74
75
76
77
78
79
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
rusty1s's avatar
rusty1s committed
80

rusty1s's avatar
rusty1s committed
81
### Graclus
rusty1s's avatar
rusty1s committed
82
83

A greedy clustering algorithm of picking an unmarked vertex and matching it with one its unmarked neighbors (that maximizes its edge weight).
rusty1s's avatar
credit  
rusty1s committed
84
The GPU algorithm is adapted from Fagginger Auer and Bisseling: [A GPU Algorithm for Greedy Graph Matching](http://www.staff.science.uu.nl/~bisse101/Articles/match12.pdf) (LNCS 2012)
rusty1s's avatar
rusty1s committed
85
86
87
88
89

```python
import torch
from torch_cluster import graclus_cluster

rusty1s's avatar
rusty1s committed
90
91
row = torch.tensor([0, 1, 1, 2])
col = torch.tensor([1, 0, 2, 1])
rusty1s's avatar
new api  
rusty1s committed
92
weight = torch.Tensor([1, 1, 1, 1])  # Optional edge weights.
rusty1s's avatar
rusty1s committed
93
94
95
96
97
98

cluster = graclus_cluster(row, col, weight)
```

```
print(cluster)
rusty1s's avatar
rusty1s committed
99
tensor([0, 0, 1])
rusty1s's avatar
rusty1s committed
100
101
```

rusty1s's avatar
rusty1s committed
102
### VoxelGrid
rusty1s's avatar
rusty1s committed
103
104

A clustering algorithm, which overlays a regular grid of user-defined size over a point cloud and clusters all points within a voxel.
rusty1s's avatar
rusty1s committed
105
106
107
108
109

```python
import torch
from torch_cluster import grid_cluster

rusty1s's avatar
new api  
rusty1s committed
110
111
pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
size = torch.Tensor([5, 5])
rusty1s's avatar
rusty1s committed
112
113
114
115
116
117

cluster = grid_cluster(pos, size)
```

```
print(cluster)
rusty1s's avatar
rusty1s committed
118
tensor([0, 5, 3, 0, 1])
rusty1s's avatar
rusty1s committed
119
```
rusty1s's avatar
rusty1s committed
120

rusty1s's avatar
rusty1s committed
121
### FarthestPointSampling
rusty1s's avatar
rusty1s committed
122

rusty1s's avatar
rusty1s committed
123
A sampling algorithm, which iteratively samples the most distant point with regard to the rest points.
rusty1s's avatar
rusty1s committed
124
125
126
127
128
129
130

```python
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])
rusty1s's avatar
rusty1s committed
131
index = fps(x, batch, ratio=0.5, random_start=False)
rusty1s's avatar
rusty1s committed
132
133
134
135
136
137
138
```

```
print(sample)
tensor([0, 3])
```

rusty1s's avatar
rusty1s committed
139
### kNN-Graph
rusty1s's avatar
rusty1s committed
140

rusty1s's avatar
rusty1s committed
141
Computes graph edges to the nearest *k* points.
rusty1s's avatar
rusty1s committed
142
143
144
145
146
147
148
149
150
151
152
153

```python
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)
Matthias Fey's avatar
Matthias Fey committed
154
155
tensor([[1, 2, 0, 3, 0, 3, 1, 2],
        [0, 0, 1, 1, 2, 2, 3, 3]])
rusty1s's avatar
rusty1s committed
156
157
```

rusty1s's avatar
rusty1s committed
158
### Radius-Graph
rusty1s's avatar
rusty1s committed
159

rusty1s's avatar
rusty1s committed
160
Computes graph edges to all points within a given distance.
rusty1s's avatar
rusty1s committed
161
162
163
164
165
166
167
168
169
170
171
172

```python
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=1.5, batch=batch, loop=False)
```

```
print(edge_index)
Matthias Fey's avatar
Matthias Fey committed
173
174
tensor([[1, 2, 0, 3, 0, 3, 1, 2],
        [0, 0, 1, 1, 2, 2, 3, 3]])
rusty1s's avatar
rusty1s committed
175
176
```

rusty1s's avatar
rusty1s committed
177
### Nearest
rusty1s's avatar
rusty1s committed
178

rusty1s's avatar
rusty1s committed
179
Clusters points in *x* together which are nearest to a given query point in *y*.
rusty1s's avatar
rusty1s committed
180
181
182
183
184
185

```python
import torch
from torch_cluster import nearest

x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
rusty1s's avatar
rusty1s committed
186
187
188
189
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)
rusty1s's avatar
rusty1s committed
190
191
192
193
194
195
196
```

```
print(cluster)
tensor([0, 0, 1, 1])
```

rusty1s's avatar
rusty1s committed
197
### RandomWalk-Sampling
rusty1s's avatar
rusty1s committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213

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

```python
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)
Maurits's avatar
Maurits committed
214
tensor([[0, 1, 2, 4],
rusty1s's avatar
rusty1s committed
215
        [1, 3, 4, 2],
Maurits's avatar
Maurits committed
216
217
218
        [2, 4, 2, 1],
        [3, 4, 2, 4],
        [4, 3, 1, 0]])
rusty1s's avatar
rusty1s committed
219
220
```

rusty1s's avatar
rusty1s committed
221
222
223
224
225
## Running tests

```
python setup.py test
```