README.md 7.08 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
linting  
rusty1s committed
16
[PyTorch](http://pytorch.org/) completely lacks autograd support and operations such as sparse sparse matrix multiplication, but is heavily working on improvement (*cf.* [this issue](https://github.com/pytorch/pytorch/issues/9674)).
rusty1s's avatar
rusty1s committed
17
In the meantime, this package consists of a small extension library of optimized sparse matrix operations with autograd support.
rusty1s's avatar
typos  
rusty1s committed
18
This package currently consists of the following methods:
rusty1s's avatar
rusty1s committed
19

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

All included operations work on varying data types and are implemented both for CPU and GPU.
rusty1s's avatar
rusty1s committed
26
27
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
28
29
30

## Installation

rusty1s's avatar
rusty1s committed
31
Ensure that at least PyTorch 1.1.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
32
33

```
rusty1s's avatar
rusty1s committed
34
$ python -c "import torch; print(torch.__version__)"
rusty1s's avatar
typo  
rusty1s committed
35
>>> 1.1.0
rusty1s's avatar
rusty1s committed
36

rusty1s's avatar
rusty1s committed
37
$ echo $PATH
rusty1s's avatar
rusty1s committed
38
>>> /usr/local/cuda/bin:...
rusty1s's avatar
rusty1s committed
39
40

$ echo $CPATH
rusty1s's avatar
rusty1s committed
41
>>> /usr/local/cuda/include:...
rusty1s's avatar
rusty1s committed
42
43
44
45
46
```

Then run:

```
rusty1s's avatar
rusty1s committed
47
pip install torch-scatter torch-sparse
rusty1s's avatar
rusty1s committed
48
49
```

rusty1s's avatar
cleanup  
rusty1s committed
50
If you are running into any installation problems, please create an [issue](https://github.com/rusty1s/pytorch_sparse/issues).
rusty1s's avatar
rusty1s committed
51
Be sure to import `torch` first before using this package to resolve symbols the dynamic linker must see.
rusty1s's avatar
links  
rusty1s committed
52

rusty1s's avatar
rusty1s committed
53
## Coalesce
rusty1s's avatar
rusty1s committed
54

rusty1s's avatar
docs  
rusty1s committed
55
```
rusty1s's avatar
rusty1s committed
56
torch_sparse.coalesce(index, value, m, n, op="add", fill_value=0) -> (torch.LongTensor, torch.Tensor)
rusty1s's avatar
docs  
rusty1s committed
57
58
```

59
Row-wise sorts `index` and removes duplicate entries.
rusty1s's avatar
rusty1s committed
60
61
62
63
64
65
66
Duplicate entries are removed by scattering them together.
For scattering, any operation of [`torch_scatter`](https://github.com/rusty1s/pytorch_scatter) can be used.

### Parameters

* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
67
68
* **m** *(int)* - The first dimension of corresponding dense matrix.
* **n** *(int)* - The second dimension of corresponding dense matrix.
rusty1s's avatar
docs  
rusty1s committed
69
70
* **op** *(string, optional)* - The scatter operation to use. (default: `"add"`)
* **fill_value** *(int, optional)* - The initial fill value of scatter operation. (default: `0`)
rusty1s's avatar
rusty1s committed
71
72
73

### Returns

rusty1s's avatar
docs  
rusty1s committed
74
75
* **index** *(LongTensor)* - The coalesced index tensor of sparse matrix.
* **value** *(Tensor)* - The coalesced value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
76
77

### Example
rusty1s's avatar
docs  
rusty1s committed
78
79

```python
ekka's avatar
ekka committed
80
import torch
rusty1s's avatar
rusty1s committed
81
82
83
84
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
85
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
rusty1s's avatar
docs  
rusty1s committed
86

rusty1s's avatar
rusty1s committed
87
index, value = coalesce(index, value, m=3, n=2)
rusty1s's avatar
docs  
rusty1s committed
88
89
```

rusty1s's avatar
rusty1s committed
90
91
92
93
94
```
print(index)
tensor([[0, 1, 1, 2],
        [1, 0, 1, 0]])
print(value)
rusty1s's avatar
rusty1s committed
95
96
97
98
tensor([[6.0, 8.0],
        [7.0, 9.0],
        [3.0, 4.0],
        [5.0, 6.0]])
rusty1s's avatar
rusty1s committed
99
```
rusty1s's avatar
docs  
rusty1s committed
100

rusty1s's avatar
rusty1s committed
101
## Transpose
rusty1s's avatar
rusty1s committed
102

rusty1s's avatar
docs  
rusty1s committed
103
```
rusty1s's avatar
rusty1s committed
104
torch_sparse.transpose(index, value, m, n) -> (torch.LongTensor, torch.Tensor)
rusty1s's avatar
docs  
rusty1s committed
105
106
```

rusty1s's avatar
rusty1s committed
107
108
109
110
111
112
Transposes dimensions 0 and 1 of a sparse matrix.

### Parameters

* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
113
114
* **m** *(int)* - The first dimension of corresponding dense matrix.
* **n** *(int)* - The second dimension of corresponding dense matrix.
rusty1s's avatar
rusty1s committed
115
116
117

### Returns

rusty1s's avatar
docs  
rusty1s committed
118
119
* **index** *(LongTensor)* - The transposed index tensor of sparse matrix.
* **value** *(Tensor)* - The transposed value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
120
121

### Example
rusty1s's avatar
docs  
rusty1s committed
122
123

```python
ekka's avatar
ekka committed
124
import torch
rusty1s's avatar
rusty1s committed
125
126
127
128
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
129
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
rusty1s's avatar
docs  
rusty1s committed
130

rusty1s's avatar
docs  
rusty1s committed
131
index, value = transpose(index, value, 3, 2)
rusty1s's avatar
docs  
rusty1s committed
132
133
```

rusty1s's avatar
rusty1s committed
134
135
136
137
138
```
print(index)
tensor([[0, 0, 1, 1],
        [1, 2, 0, 1]])
print(value)
rusty1s's avatar
rusty1s committed
139
140
141
142
tensor([[7.0, 9.0],
        [5.0, 6.0],
        [6.0, 8.0],
        [3.0, 4.0]])
rusty1s's avatar
rusty1s committed
143
```
rusty1s's avatar
docs  
rusty1s committed
144

rusty1s's avatar
rusty1s committed
145
## Sparse Dense Matrix Multiplication
rusty1s's avatar
rusty1s committed
146

rusty1s's avatar
docs  
rusty1s committed
147
```
148
torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor
rusty1s's avatar
docs  
rusty1s committed
149
150
```

rusty1s's avatar
rusty1s committed
151
152
153
Matrix product of a sparse matrix with a dense matrix.

### Parameters
rusty1s's avatar
docs  
rusty1s committed
154

rusty1s's avatar
rusty1s committed
155
156
* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
157
158
* **m** *(int)* - The first dimension of corresponding dense matrix.
* **n** *(int)* - The second dimension of corresponding dense matrix.
rusty1s's avatar
docs  
rusty1s committed
159
* **matrix** *(Tensor)* - The dense matrix.
rusty1s's avatar
rusty1s committed
160
161
162

### Returns

rusty1s's avatar
docs  
rusty1s committed
163
* **out** *(Tensor)* - The dense output matrix.
rusty1s's avatar
rusty1s committed
164
165
166
167

### Example

```python
ekka's avatar
ekka committed
168
import torch
rusty1s's avatar
rusty1s committed
169
170
171
172
from torch_sparse import spmm

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

rusty1s's avatar
rusty1s committed
176
out = spmm(index, value, 3, 3, matrix)
rusty1s's avatar
rusty1s committed
177
178
179
180
```

```
print(out)
181
182
183
tensor([[7.0, 16.0],
        [8.0, 20.0],
        [7.0, 19.0]])
rusty1s's avatar
docs  
rusty1s committed
184
```
rusty1s's avatar
rusty1s committed
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200

## Sparse Sparse Matrix Multiplication

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

Matrix product of two sparse tensors.
Both input sparse matrices need to be **coalesced**.

### Parameters

* **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.
201
202
203
* **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
rusty1s committed
204
205
206

### Returns

rusty1s's avatar
docs  
rusty1s committed
207
208
* **index** *(LongTensor)* - The output index tensor of sparse matrix.
* **value** *(Tensor)* - The output value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
209
210
211
212

### Example

```python
ekka's avatar
ekka committed
213
import torch
rusty1s's avatar
docs  
rusty1s committed
214
215
from torch_sparse import spspmm

rusty1s's avatar
rusty1s committed
216
indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]])
rusty1s's avatar
rusty1s committed
217
valueA = torch.Tensor([1, 2, 3, 4, 5])
rusty1s's avatar
rusty1s committed
218
219

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

rusty1s's avatar
rusty1s committed
222
223
224
225
indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2)
```

```
ekka's avatar
ekka committed
226
print(indexC)
rusty1s's avatar
rusty1s committed
227
228
tensor([[0, 1, 2],
        [0, 1, 1]])
ekka's avatar
ekka committed
229
print(valueC)
230
tensor([8.0, 6.0, 8.0])
rusty1s's avatar
docs  
rusty1s committed
231
232
```

rusty1s's avatar
rusty1s committed
233
234
235
236
237
## Running tests

```
python setup.py test
```