README.md 6.86 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.0.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
rusty1s committed
35
>>> 1.0.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
```

rusty1s's avatar
rusty1s committed
59
60
61
62
63
64
65
66
Row-wise sorts `value` and removes duplicate entries.
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.
rusty1s's avatar
docs  
rusty1s committed
67
68
69
70
* **m** *(int)* - The first dimension of sparse matrix.
* **n** *(int)* - The second dimension of sparse matrix.
* **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
rusty1s's avatar
rusty1s committed
80
81
82
83
84
from torch_sparse import coalesce

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

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

rusty1s's avatar
rusty1s committed
89
90
91
92
93
94
95
```
print(index)
tensor([[0, 1, 1, 2],
        [1, 0, 1, 0]])
print(value)
tensor([[6, 8], [7, 9], [3, 4], [5, 6]])
```
rusty1s's avatar
docs  
rusty1s committed
96

rusty1s's avatar
rusty1s committed
97
## Transpose
rusty1s's avatar
rusty1s committed
98

rusty1s's avatar
docs  
rusty1s committed
99
```
rusty1s's avatar
rusty1s committed
100
torch_sparse.transpose(index, value, m, n) -> (torch.LongTensor, torch.Tensor)
rusty1s's avatar
docs  
rusty1s committed
101
102
```

rusty1s's avatar
rusty1s committed
103
104
105
106
107
108
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.
rusty1s's avatar
docs  
rusty1s committed
109
110
* **m** *(int)* - The first dimension of sparse matrix.
* **n** *(int)* - The second dimension of sparse matrix.
rusty1s's avatar
rusty1s committed
111
112
113

### Returns

rusty1s's avatar
docs  
rusty1s committed
114
115
* **index** *(LongTensor)* - The transposed index tensor of sparse matrix.
* **value** *(Tensor)* - The transposed value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
116
117

### Example
rusty1s's avatar
docs  
rusty1s committed
118
119

```python
rusty1s's avatar
rusty1s committed
120
121
122
123
124
from torch_sparse import transpose

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

rusty1s's avatar
docs  
rusty1s committed
126
index, value = transpose(index, value, 3, 2)
rusty1s's avatar
docs  
rusty1s committed
127
128
```

rusty1s's avatar
rusty1s committed
129
130
131
132
133
134
135
136
137
138
```
print(index)
tensor([[0, 0, 1, 1],
        [1, 2, 0, 1]])
print(value)
tensor([[7, 9],
        [5, 6],
        [6, 8],
        [3, 4]])
```
rusty1s's avatar
docs  
rusty1s committed
139

rusty1s's avatar
rusty1s committed
140
## Sparse Dense Matrix Multiplication
rusty1s's avatar
rusty1s committed
141

rusty1s's avatar
docs  
rusty1s committed
142
```
rusty1s's avatar
rusty1s committed
143
torch_sparse.spmm(index, value, m, matrix) -> torch.Tensor
rusty1s's avatar
docs  
rusty1s committed
144
145
```

rusty1s's avatar
rusty1s committed
146
147
148
Matrix product of a sparse matrix with a dense matrix.

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

rusty1s's avatar
rusty1s committed
150
151
* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
rusty1s's avatar
docs  
rusty1s committed
152
153
* **m** *(int)* - The first dimension of sparse matrix.
* **matrix** *(Tensor)* - The dense matrix.
rusty1s's avatar
rusty1s committed
154
155
156

### Returns

rusty1s's avatar
docs  
rusty1s committed
157
* **out** *(Tensor)* - The dense output matrix.
rusty1s's avatar
rusty1s committed
158
159
160
161
162
163
164
165

### Example

```python
from torch_sparse import spmm

index = torch.tensor([[0, 0, 1, 2, 2],
                      [0, 2, 1, 0, 1]])
166
167
value = torch.tensor([1, 2, 4, 1, 3], dtype=torch.float)
matrix = torch.tensor([[1, 4], [2, 5], [3, 6]], dtype=torch.float)
rusty1s's avatar
rusty1s committed
168
169
170
171
172
173

out = spmm(index, value, 3, matrix)
```

```
print(out)
174
175
176
tensor([[7.0, 16.0],
        [8.0, 20.0],
        [7.0, 19.0]])
rusty1s's avatar
docs  
rusty1s committed
177
```
rusty1s's avatar
rusty1s committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193

## 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.
rusty1s's avatar
docs  
rusty1s committed
194
195
196
* **m** *(int)* - The first dimension of first sparse matrix.
* **k** *(int)* - The second dimension of first sparse matrix and first dimension of second sparse matrix.
* **n** *(int)* - The second dimension of second sparse matrix.
rusty1s's avatar
rusty1s committed
197
198
199

### Returns

rusty1s's avatar
docs  
rusty1s committed
200
201
* **index** *(LongTensor)* - The output index tensor of sparse matrix.
* **value** *(Tensor)* - The output value tensor of sparse matrix.
rusty1s's avatar
rusty1s committed
202
203
204
205

### Example

```python
rusty1s's avatar
docs  
rusty1s committed
206
207
from torch_sparse import spspmm

rusty1s's avatar
rusty1s committed
208
indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]])
209
valueA = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float)
rusty1s's avatar
rusty1s committed
210
211

indexB = torch.tensor([[0, 2], [1, 0]])
212
valueB = torch.tensor([2, 4], dtype=torch.float)
rusty1s's avatar
docs  
rusty1s committed
213

rusty1s's avatar
rusty1s committed
214
215
216
217
218
219
220
221
indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2)
```

```
print(index)
tensor([[0, 1, 2],
        [0, 1, 1]])
print(value)
222
tensor([8.0, 6.0, 8.0])
rusty1s's avatar
docs  
rusty1s committed
223
224
```

rusty1s's avatar
rusty1s committed
225
226
227
228
229
## Running tests

```
python setup.py test
```