README.md 5.98 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
rusty1s committed
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
### Binaries

We provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://pytorch-geometric.com/whl).
To install the binaries for PyTorch 1.4.0, simply run

```
pip install torch-spline-conv==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.4.0.html
```

where `${CUDA}` should be replaced by either `cpu`, `cu92`, `cu100` or `cu101` depending on your PyTorch installation.

|             | `cpu` | `cu92` | `cu100` | `cu101` |
|-------------|-------|--------|---------|---------|
| **Linux**   | ✅    | ✅     | ✅      | ✅      |
| **Windows** | ✅    | ❌     | ❌      | ✅      |
| **macOS**   | ✅    |        |         |         |

### From source

Ensure that at least PyTorch 1.4.0 is installed and verify that `cuda/bin` and `cuda/include` are in your `$PATH` and `$CPATH` respectively, *e.g.*:
rusty1s's avatar
rusty1s committed
44
45
46

```
$ python -c "import torch; print(torch.__version__)"
rusty1s's avatar
rusty1s committed
47
>>> 1.4.0
rusty1s's avatar
rusty1s committed
48
49
50
51
52

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

$ echo $CPATH
rusty1s's avatar
rusty1s committed
53
>>> /usr/local/cuda/include:...
rusty1s's avatar
rusty1s committed
54
55
```

rusty1s's avatar
bugfix  
rusty1s committed
56
57
58
Then run:

```
rusty1s's avatar
rusty1s committed
59
pip install torch-spline-conv
rusty1s's avatar
bugfix  
rusty1s committed
60
61
```

rusty1s's avatar
rusty1s committed
62
63
64
65
66
67
When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail.
In this case, ensure that the compute capabilities are set via `TORCH_CUDA_ARCH_LIST`, *e.g.*:

```
export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"
```
rusty1s's avatar
rusty1s committed
68

rusty1s's avatar
bugfix  
rusty1s committed
69
70
71
## Usage

```python
rusty1s's avatar
rusty1s committed
72
73
74
75
76
77
78
79
80
81
82
83
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)
rusty1s's avatar
bugfix  
rusty1s committed
84
85
```

rusty1s's avatar
typo  
rusty1s committed
86
Applies the spline-based convolution operator
rusty1s's avatar
rusty1s committed
87
<p align="center">
Matthias Fey's avatar
Matthias Fey committed
88
  <img width="50%" src="https://user-images.githubusercontent.com/6945922/38684093-36d9c52e-3e6f-11e8-9021-db054223c6b9.png" />
rusty1s's avatar
rusty1s committed
89
</p>
rusty1s's avatar
bugfix  
rusty1s committed
90
over several node features of an input graph.
rusty1s's avatar
typo  
rusty1s committed
91
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
92

Matthias Fey's avatar
Matthias Fey committed
93
94
95
96
97
<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
98
99
### Parameters

rusty1s's avatar
rusty1s committed
100
* **x** *(Tensor)* - Input node features of shape `(number_of_nodes x in_channels)`.
rusty1s's avatar
rusty1s committed
101
102
103
104
105
* **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
106
* **degree** *(int, optional)* - B-spline basis degree. (default: `1`)
rusty1s's avatar
rusty1s committed
107
* **norm** *(bool, optional)*: Whether to normalize output by node degree. (default: `True`)
rusty1s's avatar
rusty1s committed
108
109
* **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
110
111
112

### Returns

Matthias Fey's avatar
Matthias Fey committed
113
* **out** *(Tensor)* - Out node features of shape `(number_of_nodes x out_channels)`.
rusty1s's avatar
bugfix  
rusty1s committed
114
115
116
117
118

### Example

```python
import torch
rusty1s's avatar
rusty1s committed
119
from torch_spline_conv import spline_conv
rusty1s's avatar
bugfix  
rusty1s committed
120

rusty1s's avatar
rusty1s committed
121
x = torch.rand((4, 2), dtype=torch.float)  # 4 nodes with 2 features each
rusty1s's avatar
rusty1s committed
122
123
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
124
125
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
126
is_open_spline = torch.tensor([1, 1], dtype=torch.uint8)  # only use open B-splines
rusty1s's avatar
rusty1s committed
127
degree = 1  # B-spline degree of 1
rusty1s's avatar
rusty1s committed
128
norm = True  # Normalize output by node degree.
rusty1s's avatar
rusty1s committed
129
root_weight = torch.rand((2, 4), dtype=torch.float)  # separately weight root nodes
rusty1s's avatar
typo  
rusty1s committed
130
bias = None  # do not apply an additional bias
rusty1s's avatar
bugfix  
rusty1s committed
131

rusty1s's avatar
rusty1s committed
132
133
out = spline_conv(x, edge_index, pseudo, weight, kernel_size,
                  is_open_spline, degree, norm, root_weight, bias)
rusty1s's avatar
bugfix  
rusty1s committed
134

rusty1s's avatar
rename  
rusty1s committed
135
print(out.size())
rusty1s's avatar
typo  
rusty1s committed
136
torch.Size([4, 4])  # 4 nodes with 4 features each
rusty1s's avatar
bugfix  
rusty1s committed
137
138
```

rusty1s's avatar
rusty1s committed
139
140
141
142
143
144
145
146
## 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},
Matthias Fey's avatar
Matthias Fey committed
147
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
rusty1s's avatar
rusty1s committed
148
149
150
  year={2018},
}
```
rusty1s's avatar
typos  
rusty1s committed
151
152
153
154
155
156

## Running tests

```
python setup.py test
```