README.md 4.31 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
[pypi-image]: https://badge.fury.io/py/torch-scatter.svg
[pypi-url]: https://pypi.python.org/pypi/torch-scatter
rusty1s's avatar
rusty1s committed
3
4
[build-image]: https://travis-ci.org/rusty1s/pytorch_scatter.svg?branch=master
[build-url]: https://travis-ci.org/rusty1s/pytorch_scatter
rusty1s's avatar
rusty1s committed
5
6
[docs-image]: https://readthedocs.org/projects/pytorch-scatter/badge/?version=latest
[docs-url]: https://pytorch-scatter.readthedocs.io/en/latest/?badge=latest
rusty1s's avatar
rusty1s committed
7
8
[coverage-image]: https://codecov.io/gh/rusty1s/pytorch_scatter/branch/master/graph/badge.svg
[coverage-url]: https://codecov.io/github/rusty1s/pytorch_scatter?branch=master
rusty1s's avatar
rusty1s committed
9

rusty1s's avatar
rusty1s committed
10
11
12
13
# PyTorch Scatter

[![PyPI Version][pypi-image]][pypi-url]
[![Build Status][build-image]][build-url]
rusty1s's avatar
rusty1s committed
14
[![Docs Status][docs-image]][docs-url]
rusty1s's avatar
rusty1s committed
15
[![Code Coverage][coverage-image]][coverage-url]
rusty1s's avatar
rusty1s committed
16

rusty1s's avatar
rusty1s committed
17
<p align="center">
rusty1s's avatar
smaller  
rusty1s committed
18
  <img width="50%" src="https://raw.githubusercontent.com/rusty1s/pytorch_scatter/master/docs/source/_figures/add.svg?sanitize=true" />
rusty1s's avatar
rusty1s committed
19
20
21
22
</p>

--------------------------------------------------------------------------------

rusty1s's avatar
rusty1s committed
23
**[Documentation](https://pytorch-scatter.readthedocs.io)**
rusty1s's avatar
rusty1s committed
24

rusty1s's avatar
rusty1s committed
25
This package consists of a small extension library of highly optimized sparse update (scatter) operations for the use in [PyTorch](http://pytorch.org/), which are missing in the main package.
rusty1s's avatar
typo  
rusty1s committed
26
Scatter operations can be roughly described as reduce operations based on a given "group-index" tensor.
rusty1s's avatar
rusty1s committed
27
The package consists of the following operations:
rusty1s's avatar
rusty1s committed
28

rusty1s's avatar
rusty1s committed
29
30
31
32
33
34
35
36
* [**Scatter Add**](https://pytorch-scatter.readthedocs.io/en/latest/functions/add.html)
* [**Scatter Sub**](https://pytorch-scatter.readthedocs.io/en/latest/functions/sub.html)
* [**Scatter Mul**](https://pytorch-scatter.readthedocs.io/en/latest/functions/mul.html)
* [**Scatter Div**](https://pytorch-scatter.readthedocs.io/en/latest/functions/div.html)
* [**Scatter Mean**](https://pytorch-scatter.readthedocs.io/en/latest/functions/mean.html)
* [**Scatter Std**](https://pytorch-scatter.readthedocs.io/en/latest/functions/std.html)
* [**Scatter Min**](https://pytorch-scatter.readthedocs.io/en/latest/functions/min.html)
* [**Scatter Max**](https://pytorch-scatter.readthedocs.io/en/latest/functions/max.html)
37
38
39
40
41
42
* [**Scatter LogSumExp**](https://pytorch-scatter.readthedocs.io/en/latest/functions/logsumexp.html)

In addition, we provide composite functions which make use of `scatter_*` operations under the hood:

* [**Scatter Softmax**](https://pytorch-scatter.readthedocs.io/en/latest/composite/softmax.html#torch_scatter.composite.scatter_softmax)
* [**Scatter LogSoftmax**](https://pytorch-scatter.readthedocs.io/en/latest/composite/softmax.html#torch_scatter.composite.scatter_log_softmax)
rusty1s's avatar
rusty1s committed
43

rusty1s's avatar
rusty1s committed
44
All included operations are broadcastable, work on varying data types, and are implemented both for CPU and GPU with corresponding backward implementations.
rusty1s's avatar
rusty1s committed
45

rusty1s's avatar
rusty1s committed
46
47
## Installation

rusty1s's avatar
rusty1s committed
48
Ensure that at least PyTorch 1.3.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
49
50
51

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

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

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

rusty1s's avatar
typos  
rusty1s committed
61
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`
rusty1s's avatar
rusty1s committed
62

rusty1s's avatar
rusty1s committed
63
```
rusty1s's avatar
typos  
rusty1s committed
64
export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"
rusty1s's avatar
rusty1s committed
65
66
```

67
68
69
70
71
### Windows

If you are installing this on Windows specifically, **you will need to point the setup to your Visual Studio installation** for some neccessary libraries and header files.
To do this, add the include and library paths of your installation to the path lists in setup.py as described in the respective comments in the code.

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

rusty1s's avatar
rusty1s committed
75
## Example
rusty1s's avatar
rusty1s committed
76

rusty1s's avatar
rusty1s committed
77
```py
rusty1s's avatar
typo  
rusty1s committed
78
import torch
rusty1s's avatar
rusty1s committed
79
80
from torch_scatter import scatter_max

rusty1s's avatar
rusty1s committed
81
82
src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
rusty1s's avatar
rusty1s committed
83

rusty1s's avatar
rusty1s committed
84
out, argmax = scatter_max(src, index, fill_value=0)
rusty1s's avatar
typo  
rusty1s committed
85
```
rusty1s's avatar
rusty1s committed
86

rusty1s's avatar
typo  
rusty1s committed
87
```
rusty1s's avatar
rusty1s committed
88
89
90
print(out)
tensor([[ 0,  0,  4,  3,  2,  0],
        [ 2,  4,  3,  0,  0,  0]])
rusty1s's avatar
rusty1s committed
91

rusty1s's avatar
typos  
rusty1s committed
92
print(argmax)
rusty1s's avatar
rusty1s committed
93
94
tensor([[-1, -1,  3,  4,  0,  1]
        [ 1,  4,  3, -1, -1, -1]])
rusty1s's avatar
typos  
rusty1s committed
95
```
rusty1s's avatar
rusty1s committed
96
97
98

## Running tests

rusty1s's avatar
rusty1s committed
99
```
rusty1s's avatar
rusty1s committed
100
101
python setup.py test
```