README.md 8.04 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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
### Binaries

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

```
pip install torch-scatter==latest+${CUDA} -f http://pytorch-scatter.s3-website.eu-central-1.amazonaws.com/whl/torch-1.4.0.html --trusted-host pytorch-scatter.s3-website.eu-central-1.amazonaws.com
pip install torch-sparse==latest+${CUDA} -f http://pytorch-sparse.s3-website.eu-central-1.amazonaws.com/whl/torch-1.4.0.html --trusted-host pytorch-sparse.s3-website.eu-central-1.amazonaws.com
```

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

rusty1s's avatar
rusty1s committed
50
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
51
52

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

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

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

Then run:

```
rusty1s's avatar
rusty1s committed
66
pip install torch-scatter torch-sparse
rusty1s's avatar
rusty1s committed
67
68
```

rusty1s's avatar
rusty1s committed
69
70
71
72
73
74
75
76
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
links  
rusty1s committed
77

rusty1s's avatar
rusty1s committed
78
### Coalesce
rusty1s's avatar
rusty1s committed
79

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

84
Row-wise sorts `index` and removes duplicate entries.
rusty1s's avatar
rusty1s committed
85
86
87
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
88
#### Parameters
rusty1s's avatar
rusty1s committed
89
90
91

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

rusty1s's avatar
rusty1s committed
96
#### Returns
rusty1s's avatar
rusty1s committed
97

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

rusty1s's avatar
rusty1s committed
101
#### Example
rusty1s's avatar
docs  
rusty1s committed
102
103

```python
ekka's avatar
ekka committed
104
import torch
rusty1s's avatar
rusty1s committed
105
106
107
108
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
109
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
rusty1s's avatar
docs  
rusty1s committed
110

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

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

rusty1s's avatar
rusty1s committed
125
### Transpose
rusty1s's avatar
rusty1s committed
126

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

rusty1s's avatar
rusty1s committed
131
132
Transposes dimensions 0 and 1 of a sparse matrix.

rusty1s's avatar
rusty1s committed
133
#### Parameters
rusty1s's avatar
rusty1s committed
134
135
136

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

rusty1s's avatar
rusty1s committed
141
#### Returns
rusty1s's avatar
rusty1s committed
142

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

rusty1s's avatar
rusty1s committed
146
#### Example
rusty1s's avatar
docs  
rusty1s committed
147
148

```python
ekka's avatar
ekka committed
149
import torch
rusty1s's avatar
rusty1s committed
150
151
152
153
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
154
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
rusty1s's avatar
docs  
rusty1s committed
155

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

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

rusty1s's avatar
rusty1s committed
170
### Sparse Dense Matrix Multiplication
rusty1s's avatar
rusty1s committed
171

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

rusty1s's avatar
rusty1s committed
176
177
Matrix product of a sparse matrix with a dense matrix.

rusty1s's avatar
rusty1s committed
178
#### Parameters
rusty1s's avatar
docs  
rusty1s committed
179

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

rusty1s's avatar
rusty1s committed
186
#### Returns
rusty1s's avatar
rusty1s committed
187

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

rusty1s's avatar
rusty1s committed
190
#### Example
rusty1s's avatar
rusty1s committed
191
192

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

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

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

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

rusty1s's avatar
rusty1s committed
211
### Sparse Sparse Matrix Multiplication
rusty1s's avatar
rusty1s committed
212
213
214
215
216
217

```
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
218
Both input sparse matrices need to be **coalesced** (use the `coalesced` attribute to force).
rusty1s's avatar
rusty1s committed
219

rusty1s's avatar
rusty1s committed
220
#### Parameters
rusty1s's avatar
rusty1s committed
221
222
223
224
225

* **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.
226
227
228
* **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
229
* **coalesced** *(bool, optional)*: If set to `True`, will coalesce both input sparse matrices. (default: `False`)
rusty1s's avatar
rusty1s committed
230

rusty1s's avatar
rusty1s committed
231
#### Returns
rusty1s's avatar
rusty1s committed
232

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

rusty1s's avatar
rusty1s committed
236
#### Example
rusty1s's avatar
rusty1s committed
237
238

```python
ekka's avatar
ekka committed
239
import torch
rusty1s's avatar
docs  
rusty1s committed
240
241
from torch_sparse import spspmm

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

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

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

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

rusty1s's avatar
rusty1s committed
259
260
261
262
263
## Running tests

```
python setup.py test
```