DOCSEARCHbyAlgolia
Submit Docs

PyTorch Scatter

repository·master·Indexed 23 days ago

https://github.com/rusty1s/pytorch_scatter

An optimized extension library for PyTorch providing sparse update operations, including scatter and segment reductions (sum, mean, min, max). It features core operations like scatter, segment_coo, and segment_csr, as well as composite functions such as scatter_std, scatter_logsumexp, scatter_softmax, and scatter_log_softmax. The library supports CPU and GPU, is fully traceable, and includes a C++ API.

Tokens
1.3K
Snippets
4
Records
9
Agent score
83%

What's inside PyTorch Scatter

  1. What is PyTorch Scatter?

    master

    PyTorch Scatter is an extension library providing highly optimized sparse update operations (scatter and segment) for PyTorch. These operations perform reductions based on a "group-index" tensor.

    Core Operations All operations support reduction types: "sum", "mean", "min", and "max".

    • scatter: Based on arbitrary indices.
    • segment_coo: Based on sorted indices.
    • segment_csr: Based on compressed indices via pointers.

    Composite Functions The package also provides composite functions built on top of scatter_* operations:

    • scatter_std
    • scatter_logsumexp
    • scatter_softmax
    • scatter_log_softmax

    All operations are broadcastable, support varying data types, work on both CPU and GPU (with backward implementations), and are fully traceable.

  2. What are PyTorch Scatter operations?

    master

    PyTorch Scatter is an extension library providing highly optimized sparse update operations for PyTorch. It fills gaps in the main PyTorch package by providing specialized scatter and segment operations.

    Core Concepts

    • Scatter operations: These are reduce operations based on a "group-index" tensor. They do not require the index tensor to be sorted.
    • Segment operations: These are also reduce operations based on a "group-index" tensor, but they require the index tensor to be sorted.

    Key Features

    • Broadcasting: All included operations are broadcastable.
    • Type Support: Works on varying data types.
    • Hardware Support: Implemented for both CPU and GPU.
    • Autograd: Includes corresponding backward implementations for gradient computation.
    • Traceability: All operations are fully traceable (e.g., for TorchScript).

  3. Install PyTorch Scatter via Binaries

    master

    To avoid compiling from source, you can install pre-built pip wheels. The installation command depends on your PyTorch version and CUDA configuration.

    Replace ${CUDA} in the command with one of the following based on your installation:

    • cpu
    • cu126, cu128, cu130, cu132 (depending on the PyTorch version)

    Installation Commands

    For PyTorch 2.12:

    pip install torch-scatter -f https://data.pyg.org/whl/torch-2.12.0+${CUDA}.html

    For PyTorch 2.11:

    pip install torch-scatter -f https://data.pyg.org/whl/torch-2.11.0+${CUDA}.html

    For PyTorch 2.10:

    pip install torch-scatter -f https://data.pyg.org/whl/torch-2.10.0+${CUDA}.html

    For PyTorch 2.9:

    pip install torch-scatter -f https://data.pyg.org/whl/torch-2.9.0+${CUDA}.html

    Compatibility Matrix (Examples)

    • PyTorch 2.12: Supports cpu, cu126, cu130, cu132 on Linux/Windows; cpu on macOS.
    • PyTorch 2.11/2.10/2.9: Supports cpu, cu126, cu128, cu130 on Linux/Windows; cpu on macOS.
    pip install torch-scatter -f https://data.pyg.org/whl/torch-2.12.0+${CUDA}.html

  4. Install PyTorch Scatter from Source

    master

    To install from source, ensure PyTorch 1.4.0 or newer is installed. You must also verify that your CUDA paths are correctly set in your environment variables.

    Prerequisites

    1. cuda/bin must be in your $PATH.
    2. cuda/include must be in your $CPATH.

    Installation Command

    pip install torch-scatter

    Docker/Non-NVIDIA Driver Note If running in a Docker container without an NVIDIA driver, PyTorch may fail to evaluate compute capabilities. You must manually set the compute capabilities using the TORCH_CUDA_ARCH_LIST environment variable.

    Example:

    export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"
    pip install torch-scatter

  5. Use scatter_max in PyTorch

    master

    The scatter_max function performs a max reduction based on an index tensor. It returns both the maximum values and the indices (argmax) where those maximums were found.

    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, dim=-1)
    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, dim=-1)

  8. Build the C++ API

    master

    The torch-scatter C++ API provides C++ equivalents of the Python models. To build it, you need to provide the TorchLib path to CMake.

    1. Obtain the cmake_prefix_path by running import torch; print(torch.utils.cmake_prefix_path) in Python.
    2. Use that path with the -DCMAKE_PREFIX_PATH flag during the CMake configuration step.

    Build Steps

    mkdir build
cd build
# Add -DWITH_CUDA=on support for CUDA support
cmake -DCMAKE_PREFIX_PATH="..." ..
make
make install
    cmake -DCMAKE_PREFIX_PATH="..." ..
make
make install