DOCSEARCHbyAlgolia
Submit Docs

torch-spline-conv

repository·master·Indexed 41 days ago

https://github.com/rusty1s/pytorch_spline_conv

A PyTorch implementation of the spline-based convolution operator from the SplineCNN paper for fast geometric deep learning using continuous B-spline kernels. It supports CPU and GPU execution, all floating-point data types, and provides both Python and C++ APIs. Note: This repository is deprecated in favor of pyg-lib>=0.7.0.

Tokens
1.6K
Snippets
5
Records
5
Agent score
2%

What's inside torch-spline-conv

  1. Install torch-spline-conv via binaries

    master

    You can install pre-compiled pip wheels for various PyTorch and CUDA combinations. The installation command requires specifying a find-links URL that matches your PyTorch version and CUDA installation.

    Replace ${CUDA} with one of the following based on your environment:

    • cpu (for CPU only)
    • cu126 (CUDA 12.6)
    • cu128 (CUDA 12.8)
    • cu130 (CUDA 13.0)
    • cu129 (for PyTorch 2.8 specifically)

    Note: This repository is deprecated in favor of pyg-lib>=0.7.0.

    # For PyTorch 2.10
pip install torch-spline-conv -f https://data.pyg.org/whl/torch-2.10.0+${CUDA}.html

# For PyTorch 2.9
pip install torch-spline-conv -f https://data.pyg.org/whl/torch-2.9.0+${CUDA}.html

# For PyTorch 2.8
pip install torch-spline-conv -f https://data.pyg.org/whl/torch-2.8.0+${CUDA}.html

  2. Install torch-spline-conv from source

    master

    To build from source, ensure PyTorch 1.4.0 or newer is installed. You must also ensure that cuda/bin and cuda/include are present in your $PATH and $CPATH respectively.

    If you are running in a Docker container without an NVIDIA driver, you may need to manually set the compute capabilities using the TORCH_CUDA_ARCH_LIST environment variable to prevent PyTorch from failing during evaluation.

    # Verify PyTorch version
python -c "import torch; print(torch.__version__)"

# Ensure CUDA paths are set
echo $PATH
echo $CPATH

# Install
pip install torch-spline-conv

# If in Docker without NVIDIA driver, set compute capabilities
export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"

  3. Complete spline_conv usage example

    master

    A full example demonstrating how to initialize tensors and call spline_conv to process a graph with 4 nodes and 2 input features into 4 output features.

    import torch
from torch_spline_conv import spline_conv

x = torch.rand((4, 2), dtype=torch.float)  # 4 nodes with 2 features each
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]])  # 6 edges
pseudo = torch.rand((6, 2), dtype=torch.float)  # two-dimensional edge attributes
weight = torch.rand((25, 2, 4), dtype=torch.float)  # 25 parameters for in_channels x out_channels
kernel_size = torch.tensor([5, 5])  # 5 parameters in each edge dimension
is_open_spline = torch.tensor([1, 1], dtype=torch.uint8)  # only use open B-splines
degree = 1  # B-spline degree of 1
norm = True  # Normalize output by node degree.
root_weight = torch.rand((2, 4), dtype=torch.float)  # separately weight root nodes
bias = None  # do not apply an additional bias

out = spline_conv(x, edge_index, pseudo, weight, kernel_size,
                  is_open_spline, degree, norm, root_weight, bias)

print(out.size())
# Expected output: torch.Size([4, 4])

  4. Build the C++ API

    master

    The project provides a C++ API that offers C++ equivalents of the Python models. You can build it using CMake. If you require CUDA support, ensure you pass the -DWITH_CUDA=on flag during the CMake configuration step.

    mkdir build
cd build
# Add -DWITH_CUDA=on support for the CUDA if needed
cmake ..
make
make install

  5. Use the spline_conv function

    master

    The spline_conv function applies a spline-based convolution operator over node features in a graph. The kernel is defined over a weighted B-spline tensor product basis.

    Parameters:

    • x (Tensor): Input node features of shape (number_of_nodes x in_channels).
    • edge_index (LongTensor): Graph edges (source and target indices) of shape (2 x number_of_edges).
    • pseudo (Tensor): Edge attributes (pseudo coordinates) in the interval [0, 1], shape (number_of_edges x number_of_edge_attributes).
    • weight (Tensor): Trainable weight parameters of shape (kernel_size x in_channels x out_channels).
    • kernel_size (LongTensor): Number of trainable weight parameters in each edge dimension.
    • is_open_spline (ByteTensor): Whether to use open or closed B-spline bases for each dimension.
    • degree (int, optional): B-spline basis degree. Default is 1.
    • norm (bool, optional): Whether to normalize output by node degree. Default is True.
    • root_weight (Tensor, optional): Additional shared trainable parameters for the root node, shape (in_channels x out_channels). Default is None.
    • bias (Tensor, optional): Optional bias of shape (out_channels). Default is None.

    Returns:

    • out (Tensor): Out node features of shape (number_of_nodes x out_channels).
    from torch_spline_conv import spline_conv

out = spline_conv(x, edge_index, pseudo, weight, kernel_size, is_open_spline, degree=1, norm=True, root_weight=None, bias=None)