ceades / pytorch_scatter

PyTorch Extension Library of Optimized Scatter Operations

PyTorch Scatter

---

Documentation

This package consists of a small extension library of highly optimized sparse update (scatter) operations for the use in PyTorch, which are missing in the main package. Scatter operations can be roughly described as reduce operations based on a given "group-index" tensor. The package consists of the following operations:

• Scatter Add
• Scatter Sub
• Scatter Mul
• Scatter Div
• Scatter Mean
• Scatter Std
• Scatter Min
• Scatter Max
• Scatter LogSumExp

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

• Scatter Softmax
• Scatter LogSoftmax

All included operations are broadcastable, work on varying data types, and are implemented both for CPU and GPU with corresponding backward implementations.

Installation

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.:

$ python -c "import torch; print(torch.__version__)"
>>> 1.1.0

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

$ echo $CPATH
>>> /usr/local/cuda/include:...

Running in a docker container without nvidia driver, pytorch needs to evaluate the compute capabilities and fails. Ensure in this case that the compute capabilities are set in ENV TORCH_CUDA_ARCH_LIST

ENV TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"

If you are running into any installation problems, please create an issue. Be sure to import torch first before using this package to resolve symbols the dynamic linker must see.

Example

import torch
from torch_scatter import scatter_max

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]])

out, argmax = scatter_max(src, index, fill_value=0)

print(out)
tensor([[ 0,  0,  4,  3,  2,  0],
        [ 2,  4,  3,  0,  0,  0]])

print(argmax)
tensor([[-1, -1,  3,  4,  0,  1]
        [ 1,  4,  3, -1, -1, -1]])

Running tests

python setup.py test