README.md 7.34 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
### Binaries

We provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://pytorch-geometric.com/whl).
rusty1s's avatar
rusty1s committed
27

rusty1s's avatar
rusty1s committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
#### PyTorch 1.6.0

To install the binaries for PyTorch 1.6.0, simply run

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

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

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

rusty1s's avatar
rusty1s committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#### PyTorch 1.5.0

To install the binaries for PyTorch 1.5.0, simply run

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

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

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

#### PyTorch 1.4.0

rusty1s's avatar
rusty1s committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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
79
80
81

```
$ python -c "import torch; print(torch.__version__)"
rusty1s's avatar
rusty1s committed
82
>>> 1.4.0
rusty1s's avatar
rusty1s committed
83
84
85
86
87

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

$ echo $CPATH
rusty1s's avatar
rusty1s committed
88
>>> /usr/local/cuda/include:...
rusty1s's avatar
rusty1s committed
89
90
```

rusty1s's avatar
bugfix  
rusty1s committed
91
92
93
Then run:

```
rusty1s's avatar
rusty1s committed
94
pip install torch-spline-conv
rusty1s's avatar
bugfix  
rusty1s committed
95
96
```

rusty1s's avatar
rusty1s committed
97
98
99
100
101
102
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
103

rusty1s's avatar
bugfix  
rusty1s committed
104
105
106
## Usage

```python
rusty1s's avatar
rusty1s committed
107
108
109
110
111
112
113
114
115
116
117
118
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
119
120
```

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

Matthias Fey's avatar
Matthias Fey committed
128
129
130
131
132
<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
133
134
### Parameters

rusty1s's avatar
rusty1s committed
135
* **x** *(Tensor)* - Input node features of shape `(number_of_nodes x in_channels)`.
rusty1s's avatar
rusty1s committed
136
137
138
139
140
* **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
141
* **degree** *(int, optional)* - B-spline basis degree. (default: `1`)
rusty1s's avatar
rusty1s committed
142
* **norm** *(bool, optional)*: Whether to normalize output by node degree. (default: `True`)
rusty1s's avatar
rusty1s committed
143
144
* **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
145
146
147

### Returns

Matthias Fey's avatar
Matthias Fey committed
148
* **out** *(Tensor)* - Out node features of shape `(number_of_nodes x out_channels)`.
rusty1s's avatar
bugfix  
rusty1s committed
149
150
151
152
153

### Example

```python
import torch
rusty1s's avatar
rusty1s committed
154
from torch_spline_conv import spline_conv
rusty1s's avatar
bugfix  
rusty1s committed
155

rusty1s's avatar
rusty1s committed
156
x = torch.rand((4, 2), dtype=torch.float)  # 4 nodes with 2 features each
rusty1s's avatar
rusty1s committed
157
158
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
159
160
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
161
is_open_spline = torch.tensor([1, 1], dtype=torch.uint8)  # only use open B-splines
rusty1s's avatar
rusty1s committed
162
degree = 1  # B-spline degree of 1
rusty1s's avatar
rusty1s committed
163
norm = True  # Normalize output by node degree.
rusty1s's avatar
rusty1s committed
164
root_weight = torch.rand((2, 4), dtype=torch.float)  # separately weight root nodes
rusty1s's avatar
typo  
rusty1s committed
165
bias = None  # do not apply an additional bias
rusty1s's avatar
bugfix  
rusty1s committed
166

rusty1s's avatar
rusty1s committed
167
168
out = spline_conv(x, edge_index, pseudo, weight, kernel_size,
                  is_open_spline, degree, norm, root_weight, bias)
rusty1s's avatar
bugfix  
rusty1s committed
169

rusty1s's avatar
rename  
rusty1s committed
170
print(out.size())
rusty1s's avatar
typo  
rusty1s committed
171
torch.Size([4, 4])  # 4 nodes with 4 features each
rusty1s's avatar
bugfix  
rusty1s committed
172
173
```

rusty1s's avatar
rusty1s committed
174
175
176
177
178
179
180
181
## 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
182
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
rusty1s's avatar
rusty1s committed
183
184
185
  year={2018},
}
```
rusty1s's avatar
typos  
rusty1s committed
186
187
188
189
190
191

## Running tests

```
python setup.py test
```
rusty1s's avatar
rusty1s committed
192
193
194
195
196
197
198
199
200
201
202
203
204

## C++ API

`torch-spline-conv` also offers a C++ API that contains C++ equivalent of python models.

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