README.md 4.94 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
6
7
[pypi-image]: https://badge.fury.io/py/torch-spline-conv.svg
[pypi-url]: https://pypi.python.org/pypi/torch-spline-conv
[build-image]: https://travis-ci.org/rusty1s/pytorch_spline_conv.svg?branch=master
[build-url]: https://travis-ci.org/rusty1s/pytorch_spline_conv
[coverage-image]: https://codecov.io/gh/rusty1s/pytorch_spline_conv/branch/master/graph/badge.svg
[coverage-url]: https://codecov.io/github/rusty1s/pytorch_spline_conv?branch=master

rusty1s's avatar
typos  
rusty1s committed
8
# Spline-Based Convolution Operator of SplineCNN
rusty1s's avatar
rusty1s committed
9
10
11
12
13
14

[![PyPI Version][pypi-image]][pypi-url]
[![Build Status][build-image]][build-url]
[![Code Coverage][coverage-image]][coverage-url]

--------------------------------------------------------------------------------
rusty1s's avatar
rusty1s committed
15

rusty1s's avatar
bugfix  
rusty1s committed
16
17
18
19
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](https://arxiv.org/abs/1711.08920) (CVPR 2018)

rusty1s's avatar
typos  
rusty1s committed
20
The operator works on all floating point data types and is implemented both for CPU and GPU.
rusty1s's avatar
bugfix  
rusty1s committed
21
22
23

## Installation

rusty1s's avatar
typos  
rusty1s committed
24
If cuda is available, check that `nvcc` is accessible from your terminal, e.g. by typing `nvcc --version`.
rusty1s's avatar
bugfix  
rusty1s committed
25
26
27
28
29
30
31
32
33
34
If not, add cuda (`/usr/local/cuda/bin`) to your `$PATH`.
Then run:

```
pip install cffi torch-spline-conv
```

## Usage

```python
rusty1s's avatar
rusty1s committed
35
from torch_spline_conv import SplineConv
rusty1s's avatar
bugfix  
rusty1s committed
36

rusty1s's avatar
rename  
rusty1s committed
37
out = SplineConv.apply(src,
rusty1s's avatar
typo  
rusty1s committed
38
39
40
41
42
43
                       edge_index,
                       pseudo,
                       weight,
                       kernel_size,
                       is_open_spline,
                       degree=1,
rusty1s's avatar
rusty1s committed
44
                       norm=True,
rusty1s's avatar
typo  
rusty1s committed
45
46
                       root_weight=None,
                       bias=None)
rusty1s's avatar
bugfix  
rusty1s committed
47
48
```

rusty1s's avatar
typo  
rusty1s committed
49
Applies the spline-based convolution operator
rusty1s's avatar
rusty1s committed
50
<p align="center">
Matthias Fey's avatar
Matthias Fey committed
51
  <img width="50%" src="https://user-images.githubusercontent.com/6945922/38684093-36d9c52e-3e6f-11e8-9021-db054223c6b9.png" />
rusty1s's avatar
rusty1s committed
52
</p>
rusty1s's avatar
bugfix  
rusty1s committed
53
over several node features of an input graph.
rusty1s's avatar
typo  
rusty1s committed
54
The kernel function is defined over the weighted B-spline tensor product basis, as shown below for different B-spline degrees.
rusty1s's avatar
bugfix  
rusty1s committed
55

Matthias Fey's avatar
Matthias Fey committed
56
57
58
59
60
<p align="center">
  <img width="45%" src="https://user-images.githubusercontent.com/6945922/38685443-3a2a0c68-3e72-11e8-8e13-9ce9ad8fe43e.png" />
  <img width="45%" src="https://user-images.githubusercontent.com/6945922/38685459-42b2bcae-3e72-11e8-88cc-4b61e41dbd93.png" />
</p>

rusty1s's avatar
bugfix  
rusty1s committed
61
62
### Parameters

rusty1s's avatar
rusty1s committed
63
64
65
66
67
68
* **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.
rusty1s's avatar
rusty1s committed
69
* **degree** *(int, optional)* - B-spline basis degree. (default: `1`)
rusty1s's avatar
rusty1s committed
70
* **norm** *(bool, optional)*: Whether to normalize output by node degree. (default: `True`)
rusty1s's avatar
rusty1s committed
71
72
* **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`)
rusty1s's avatar
return  
rusty1s committed
73
74
75

### Returns

rusty1s's avatar
rename  
rusty1s committed
76
* **out** *(Tensor)* - out node features of shape `(number_of_nodes x out_channels)`.
rusty1s's avatar
bugfix  
rusty1s committed
77
78
79
80
81

### Example

```python
import torch
rusty1s's avatar
rusty1s committed
82
from torch_spline_conv import SplineConv
rusty1s's avatar
bugfix  
rusty1s committed
83

rusty1s's avatar
rusty1s committed
84
85
86
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
rusty1s's avatar
typo  
rusty1s committed
87
88
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
rusty1s's avatar
rusty1s committed
89
is_open_spline = torch.tensor([1, 1], dtype=torch.uint8)  # only use open B-splines
rusty1s's avatar
rusty1s committed
90
degree = 1  # B-spline degree of 1
rusty1s's avatar
rusty1s committed
91
norm = True  # Normalize output by node degree.
rusty1s's avatar
rusty1s committed
92
root_weight = torch.rand((2, 4), dtype=torch.float)  # separately weight root nodes
rusty1s's avatar
typo  
rusty1s committed
93
bias = None  # do not apply an additional bias
rusty1s's avatar
bugfix  
rusty1s committed
94

rusty1s's avatar
rename  
rusty1s committed
95
out = SplineConv.apply(src, edge_index, pseudo, weight, kernel_size,
rusty1s's avatar
typo  
rusty1s committed
96
                       is_open_spline, degree, norm, root_weight, bias)
rusty1s's avatar
bugfix  
rusty1s committed
97

rusty1s's avatar
rename  
rusty1s committed
98
print(out.size())
rusty1s's avatar
typo  
rusty1s committed
99
torch.Size([4, 4])  # 4 nodes with 4 features each
rusty1s's avatar
bugfix  
rusty1s committed
100
101
```

rusty1s's avatar
rusty1s committed
102
103
104
105
106
107
108
109
110
111
112
113
## 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},
}
```
rusty1s's avatar
typos  
rusty1s committed
114
115
116
117
118
119

## Running tests

```
python setup.py test
```