Skip to content

Latest commit

 

History

History
119 lines (91 loc) · 4.94 KB

README.md

File metadata and controls

119 lines (91 loc) · 4.94 KB

Spline-Based Convolution Operator of SplineCNN

PyPI Version Build Status Code Coverage

This is a PyTorch implementation of the spline-based convolution operator of SplineCNN, as described in our paper:

Matthias Fey, Jan Eric Lenssen, Frank Weichert, Heinrich Müller: SplineCNN: Fast Geometric Deep Learning with Continuous B-Spline Kernels (CVPR 2018)

The operator works on all floating point data types and is implemented both for CPU and GPU.

Installation

If cuda is available, check that nvcc is accessible from your terminal, e.g. by typing nvcc --version. If not, add cuda (/usr/local/cuda/bin) to your $PATH. Then run:

pip install cffi torch-spline-conv

Usage

from torch_spline_conv import SplineConv

out = SplineConv.apply(src,
                       edge_index,
                       pseudo,
                       weight,
                       kernel_size,
                       is_open_spline,
                       degree=1,
                       norm=True,
                       root_weight=None,
                       bias=None)

Applies the spline-based convolution operator

over several node features of an input graph. The kernel function is defined over the weighted B-spline tensor product basis, as shown below for different B-spline degrees.

Parameters

  • src (Tensor) - Input node features of shape (number_of_nodes x in_channels).
  • edge_index (LongTensor) - Graph edges, given by source and target indices, of shape (2 x number_of_edges).
  • pseudo (Tensor) - Edge attributes, ie. pseudo coordinates, of shape (number_of_edges x number_of_edge_attributes) in the fixed interval [0, 1].
  • 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: 1)
  • norm (bool, optional): Whether to normalize output by node degree. (default: True)
  • root_weight (Tensor, optional) - Additional shared trainable parameters for each feature of the root node of shape (in_channels x out_channels). (default: None)
  • bias (Tensor, optional) - Optional bias of shape (out_channels). (default: None)

Returns

  • out (Tensor) - out node features of shape (number_of_nodes x out_channels).

Example

import torch
from torch_spline_conv import SplineConv

src = 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 = SplineConv.apply(src, edge_index, pseudo, weight, kernel_size,
                       is_open_spline, degree, norm, root_weight, bias)

print(out.size())
torch.Size([4, 4])  # 4 nodes with 4 features each

Cite

Please cite our paper if you use this code in your own work:

@inproceedings{Fey/etal/2018,
  title={{SplineCNN}: Fast Geometric Deep Learning with Continuous {B}-Spline Kernels},
  author={Fey, Matthias and Lenssen, Jan Eric and Weichert, Frank and M{\"u}ller, Heinrich},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2018},
}

Running tests

python setup.py test