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

rusty1s's avatar
rusty1s committed
8
# PyTorch Sparse
rusty1s's avatar
rusty1s committed
9
10
11
12
13
14

[![PyPI Version][pypi-image]][pypi-url]
[![Build Status][build-image]][build-url]
[![Code Coverage][coverage-image]][coverage-url]

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

rusty1s's avatar
rusty1s committed
16
This package consists of a small extension library of optimized sparse matrix operations with autograd support.
rusty1s's avatar
typos  
rusty1s committed
17
This package currently consists of the following methods:
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
rusty1s committed
19
20
21
* **[Coalesce](#coalesce)**
* **[Transpose](#transpose)**
* **[Sparse Dense Matrix Multiplication](#sparse-dense-matrix-multiplication)**
rusty1s's avatar
docs  
rusty1s committed
22
* **[Sparse Sparse Matrix Multiplication](#sparse-sparse-matrix-multiplication)**
rusty1s's avatar
rusty1s committed
23
24

All included operations work on varying data types and are implemented both for CPU and GPU.
rusty1s's avatar
rusty1s committed
25
26
To avoid the hazzle of creating [`torch.sparse_coo_tensor`](https://pytorch.org/docs/stable/torch.html?highlight=sparse_coo_tensor#torch.sparse_coo_tensor), this package defines operations on sparse tensors by simply passing `index` and `value` tensors as arguments ([with same shapes as defined in PyTorch](https://pytorch.org/docs/stable/sparse.html)).
Note that only `value` comes with autograd support, as `index` is discrete and therefore not differentiable.
rusty1s's avatar
rusty1s committed
27
28
29

## Installation

rusty1s's avatar
rusty1s committed
30
31
### Binaries

rusty1s's avatar
rusty1s committed
32
33
We provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://pytorch-geometric.com/whl).
To install the binaries for PyTorch 1.4.0, simply run
rusty1s's avatar
rusty1s committed
34
35

```
rusty1s's avatar
rusty1s committed
36
pip install torch-scatter==latest+${CUDA} torch-sparse==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.4.0.html
rusty1s's avatar
rusty1s committed
37
38
39
40
41
42
43
44
```

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

|             | `cpu` | `cu92` | `cu100` | `cu101` |
|-------------|-------|--------|---------|---------|
| **Linux**   | ✅    | ✅     | ✅      | ✅      |
| **Windows** | ✅    | ❌     | ❌      | ✅      |
rusty1s's avatar
rusty1s committed
45
| **macOS**   | ✅    |        |         |         |
rusty1s's avatar
rusty1s committed
46
47
48

### From source

rusty1s's avatar
rusty1s committed
49
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.*:
rusty1s's avatar
rusty1s committed
50
51

```
rusty1s's avatar
rusty1s committed
52
$ python -c "import torch; print(torch.__version__)"
rusty1s's avatar
rusty1s committed
53
>>> 1.4.0
rusty1s's avatar
rusty1s committed
54

rusty1s's avatar
rusty1s committed
55
$ echo $PATH
rusty1s's avatar
rusty1s committed
56
>>> /usr/local/cuda/bin:...
rusty1s's avatar
rusty1s committed
57
58

$ echo $CPATH
rusty1s's avatar
rusty1s committed
59
>>> /usr/local/cuda/include:...
rusty1s's avatar
rusty1s committed
60
61
```

rusty1s's avatar
rusty1s committed
62
63
If you want to additionally build `torch-sparse` with METIS support, *e.g.* for partioning, please download and install the [METIS library](http://glaros.dtc.umn.edu/gkhome/metis/metis/download) by following the instructions in the `Install.txt` file.
Afterwards, set the environment variable `WITH_METIS=1`.
rusty1s's avatar
rusty1s committed
64

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

```
rusty1s's avatar
rusty1s committed
68
pip install torch-scatter torch-sparse
rusty1s's avatar
rusty1s committed
69
70
```

rusty1s's avatar
rusty1s committed
71
When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail.
rusty1s's avatar
rusty1s committed
72
73
74
In this case, ensure that the compute capabilities are set via `TORCH_CUDA_ARCH_LIST`, *e.g.*:

```
rusty1s's avatar
rusty1s committed
75
export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.2+PTX 7.5+PTX"
rusty1s's avatar
rusty1s committed
76
77
78
```

## Functions
rusty1s's avatar
links  
rusty1s committed
79

rusty1s's avatar
rusty1s committed
80
### Coalesce
rusty1s's avatar
rusty1s committed
81

rusty1s's avatar
docs  
rusty1s committed
82
```
rusty1s's avatar
rusty1s committed
83
torch_sparse.coalesce(index, value, m, n, op="add") -> (torch.LongTensor, torch.Tensor)
rusty1s's avatar
docs  
rusty1s committed
84
85
```

86
Row-wise sorts `index` and removes duplicate entries.
rusty1s's avatar
rusty1s committed
87
88
89
Duplicate entries are removed by scattering them together.
For scattering, any operation of [`torch_scatter`](https://github.com/rusty1s/pytorch_scatter) can be used.

rusty1s's avatar
rusty1s committed
90
#### Parameters
rusty1s's avatar
rusty1s committed
91
92
93

* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
94
95
* **m** *(int)* - The first dimension of corresponding dense matrix.
* **n** *(int)* - The second dimension of corresponding dense matrix.
rusty1s's avatar
docs  
rusty1s committed
96
* **op** *(string, optional)* - The scatter operation to use. (default: `"add"`)
rusty1s's avatar
rusty1s committed
97

rusty1s's avatar
rusty1s committed
98
#### Returns
rusty1s's avatar
rusty1s committed
99

rusty1s's avatar
docs  
rusty1s committed
100
101
* **index** *(LongTensor)* - The coalesced index tensor of sparse matrix.
* **value** *(Tensor)* - The coalesced value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
102

rusty1s's avatar
rusty1s committed
103
#### Example
rusty1s's avatar
docs  
rusty1s committed
104
105

```python
ekka's avatar
ekka committed
106
import torch
rusty1s's avatar
rusty1s committed
107
108
109
110
from torch_sparse import coalesce

index = torch.tensor([[1, 0, 1, 0, 2, 1],
                      [0, 1, 1, 1, 0, 0]])
rusty1s's avatar
rusty1s committed
111
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
rusty1s's avatar
docs  
rusty1s committed
112

rusty1s's avatar
rusty1s committed
113
index, value = coalesce(index, value, m=3, n=2)
rusty1s's avatar
docs  
rusty1s committed
114
115
```

rusty1s's avatar
rusty1s committed
116
117
118
119
120
```
print(index)
tensor([[0, 1, 1, 2],
        [1, 0, 1, 0]])
print(value)
rusty1s's avatar
rusty1s committed
121
122
123
124
tensor([[6.0, 8.0],
        [7.0, 9.0],
        [3.0, 4.0],
        [5.0, 6.0]])
rusty1s's avatar
rusty1s committed
125
```
rusty1s's avatar
docs  
rusty1s committed
126

rusty1s's avatar
rusty1s committed
127
### Transpose
rusty1s's avatar
rusty1s committed
128

rusty1s's avatar
docs  
rusty1s committed
129
```
rusty1s's avatar
rusty1s committed
130
torch_sparse.transpose(index, value, m, n) -> (torch.LongTensor, torch.Tensor)
rusty1s's avatar
docs  
rusty1s committed
131
132
```

rusty1s's avatar
rusty1s committed
133
134
Transposes dimensions 0 and 1 of a sparse matrix.

rusty1s's avatar
rusty1s committed
135
#### Parameters
rusty1s's avatar
rusty1s committed
136
137
138

* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
139
140
* **m** *(int)* - The first dimension of corresponding dense matrix.
* **n** *(int)* - The second dimension of corresponding dense matrix.
rusty1s's avatar
typo  
rusty1s committed
141
* **coalesced** *(bool, optional)* - If set to `False`, will not coalesce the output. (default: `True`)
rusty1s's avatar
rusty1s committed
142

rusty1s's avatar
rusty1s committed
143
#### Returns
rusty1s's avatar
rusty1s committed
144

rusty1s's avatar
docs  
rusty1s committed
145
146
* **index** *(LongTensor)* - The transposed index tensor of sparse matrix.
* **value** *(Tensor)* - The transposed value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
147

rusty1s's avatar
rusty1s committed
148
#### Example
rusty1s's avatar
docs  
rusty1s committed
149
150

```python
ekka's avatar
ekka committed
151
import torch
rusty1s's avatar
rusty1s committed
152
153
154
155
from torch_sparse import transpose

index = torch.tensor([[1, 0, 1, 0, 2, 1],
                      [0, 1, 1, 1, 0, 0]])
rusty1s's avatar
rusty1s committed
156
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
rusty1s's avatar
docs  
rusty1s committed
157

rusty1s's avatar
docs  
rusty1s committed
158
index, value = transpose(index, value, 3, 2)
rusty1s's avatar
docs  
rusty1s committed
159
160
```

rusty1s's avatar
rusty1s committed
161
162
163
164
165
```
print(index)
tensor([[0, 0, 1, 1],
        [1, 2, 0, 1]])
print(value)
rusty1s's avatar
rusty1s committed
166
167
168
169
tensor([[7.0, 9.0],
        [5.0, 6.0],
        [6.0, 8.0],
        [3.0, 4.0]])
rusty1s's avatar
rusty1s committed
170
```
rusty1s's avatar
docs  
rusty1s committed
171

rusty1s's avatar
rusty1s committed
172
### Sparse Dense Matrix Multiplication
rusty1s's avatar
rusty1s committed
173

rusty1s's avatar
docs  
rusty1s committed
174
```
175
torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor
rusty1s's avatar
docs  
rusty1s committed
176
177
```

rusty1s's avatar
rusty1s committed
178
179
Matrix product of a sparse matrix with a dense matrix.

rusty1s's avatar
rusty1s committed
180
#### Parameters
rusty1s's avatar
docs  
rusty1s committed
181

rusty1s's avatar
rusty1s committed
182
183
* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
184
185
* **m** *(int)* - The first dimension of corresponding dense matrix.
* **n** *(int)* - The second dimension of corresponding dense matrix.
rusty1s's avatar
docs  
rusty1s committed
186
* **matrix** *(Tensor)* - The dense matrix.
rusty1s's avatar
rusty1s committed
187

rusty1s's avatar
rusty1s committed
188
#### Returns
rusty1s's avatar
rusty1s committed
189

rusty1s's avatar
docs  
rusty1s committed
190
* **out** *(Tensor)* - The dense output matrix.
rusty1s's avatar
rusty1s committed
191

rusty1s's avatar
rusty1s committed
192
#### Example
rusty1s's avatar
rusty1s committed
193
194

```python
ekka's avatar
ekka committed
195
import torch
rusty1s's avatar
rusty1s committed
196
197
198
199
from torch_sparse import spmm

index = torch.tensor([[0, 0, 1, 2, 2],
                      [0, 2, 1, 0, 1]])
rusty1s's avatar
rusty1s committed
200
201
value = torch.Tensor([1, 2, 4, 1, 3])
matrix = torch.Tensor([[1, 4], [2, 5], [3, 6]])
rusty1s's avatar
rusty1s committed
202

rusty1s's avatar
rusty1s committed
203
out = spmm(index, value, 3, 3, matrix)
rusty1s's avatar
rusty1s committed
204
205
206
207
```

```
print(out)
208
209
210
tensor([[7.0, 16.0],
        [8.0, 20.0],
        [7.0, 19.0]])
rusty1s's avatar
docs  
rusty1s committed
211
```
rusty1s's avatar
rusty1s committed
212

rusty1s's avatar
rusty1s committed
213
### Sparse Sparse Matrix Multiplication
rusty1s's avatar
rusty1s committed
214
215
216
217
218
219

```
torch_sparse.spspmm(indexA, valueA, indexB, valueB, m, k, n) -> (torch.LongTensor, torch.Tensor)
```

Matrix product of two sparse tensors.
rusty1s's avatar
typo  
rusty1s committed
220
Both input sparse matrices need to be **coalesced** (use the `coalesced` attribute to force).
rusty1s's avatar
rusty1s committed
221

rusty1s's avatar
rusty1s committed
222
#### Parameters
rusty1s's avatar
rusty1s committed
223
224
225
226
227

* **indexA** *(LongTensor)* - The index tensor of first sparse matrix.
* **valueA** *(Tensor)* - The value tensor of first sparse matrix.
* **indexB** *(LongTensor)* - The index tensor of second sparse matrix.
* **valueB** *(Tensor)* - The value tensor of second sparse matrix.
228
229
230
* **m** *(int)* - The first dimension of first corresponding dense matrix.
* **k** *(int)* - The second dimension of first corresponding dense matrix and first dimension of second corresponding dense matrix.
* **n** *(int)* - The second dimension of second corresponding dense matrix.
rusty1s's avatar
typo  
rusty1s committed
231
* **coalesced** *(bool, optional)*: If set to `True`, will coalesce both input sparse matrices. (default: `False`)
rusty1s's avatar
rusty1s committed
232

rusty1s's avatar
rusty1s committed
233
#### Returns
rusty1s's avatar
rusty1s committed
234

rusty1s's avatar
docs  
rusty1s committed
235
236
* **index** *(LongTensor)* - The output index tensor of sparse matrix.
* **value** *(Tensor)* - The output value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
237

rusty1s's avatar
rusty1s committed
238
#### Example
rusty1s's avatar
rusty1s committed
239
240

```python
ekka's avatar
ekka committed
241
import torch
rusty1s's avatar
docs  
rusty1s committed
242
243
from torch_sparse import spspmm

rusty1s's avatar
rusty1s committed
244
indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]])
rusty1s's avatar
rusty1s committed
245
valueA = torch.Tensor([1, 2, 3, 4, 5])
rusty1s's avatar
rusty1s committed
246
247

indexB = torch.tensor([[0, 2], [1, 0]])
rusty1s's avatar
rusty1s committed
248
valueB = torch.Tensor([2, 4])
rusty1s's avatar
docs  
rusty1s committed
249

rusty1s's avatar
rusty1s committed
250
251
252
253
indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2)
```

```
ekka's avatar
ekka committed
254
print(indexC)
rusty1s's avatar
rusty1s committed
255
256
tensor([[0, 1, 2],
        [0, 1, 1]])
ekka's avatar
ekka committed
257
print(valueC)
258
tensor([8.0, 6.0, 8.0])
rusty1s's avatar
docs  
rusty1s committed
259
260
```

rusty1s's avatar
rusty1s committed
261
262
263
264
265
## Running tests

```
python setup.py test
```