README.md 7.52 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
[pypi-image]: https://badge.fury.io/py/torch-spline-conv.svg
[pypi-url]: https://pypi.python.org/pypi/torch-spline-conv
3
4
5
6
[testing-image]: https://github.com/rusty1s/pytorch_spline_conv/actions/workflows/testing.yml/badge.svg
[testing-url]: https://github.com/rusty1s/pytorch_spline_conv/actions/workflows/testing.yml
[linting-image]: https://github.com/rusty1s/pytorch_spline_conv/actions/workflows/linting.yml/badge.svg
[linting-url]: https://github.com/rusty1s/pytorch_spline_conv/actions/workflows/linting.yml
rusty1s's avatar
rusty1s committed
7
8
9
[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
10
# Spline-Based Convolution Operator of SplineCNN
rusty1s's avatar
rusty1s committed
11
12

[![PyPI Version][pypi-image]][pypi-url]
13
14
[![Testing Status][testing-image]][testing-url]
[![Linting Status][linting-image]][linting-url]
rusty1s's avatar
rusty1s committed
15
16
17
[![Code Coverage][coverage-image]][coverage-url]

--------------------------------------------------------------------------------
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
bugfix  
rusty1s committed
19
20
21
22
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
23
The operator works on all floating point data types and is implemented both for CPU and GPU.
rusty1s's avatar
bugfix  
rusty1s committed
24
25
26

## Installation

rusty1s's avatar
rusty1s committed
27
28
29
30
31
32
33
34
35
### Anaconda

**Update:** You can now install `pytorch-spline-conv` via [Anaconda](https://anaconda.org/rusty1s/pytorch-spline-conv) for all major OS/PyTorch/CUDA combinations 🤗
Given that you have [`pytorch >= 1.8.0` installed](https://pytorch.org/get-started/locally/), simply run

```
conda install pytorch-spline-conv -c rusty1s
```

rusty1s's avatar
rusty1s committed
36
37
### Binaries

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

40
#### PyTorch 1.9.0
rusty1s's avatar
rusty1s committed
41

42
To install the binaries for PyTorch 1.9.0, simply run
rusty1s's avatar
rusty1s committed
43
44

```
45
pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.9.0+${CUDA}.html
rusty1s's avatar
rusty1s committed
46
47
```

48
where `${CUDA}` should be replaced by either `cpu`, `cu102`, or `cu111` depending on your PyTorch installation.
rusty1s's avatar
rusty1s committed
49

50
51
52
53
54
|             | `cpu` | `cu102` | `cu111` |
|-------------|-------|---------|---------|
| **Linux**   | ✅    | ✅      | ✅      |
| **Windows** | ✅    | ✅      | ✅      |
| **macOS**   | ✅    |         |         |
rusty1s's avatar
rusty1s committed
55

56
#### PyTorch 1.8.0/1.8.1
rusty1s's avatar
rusty1s committed
57

58
To install the binaries for PyTorch 1.8.0 and 1.8.1, simply run
rusty1s's avatar
rusty1s committed
59
60

```
61
pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
rusty1s's avatar
rusty1s committed
62
63
```

64
where `${CUDA}` should be replaced by either `cpu`, `cu101`, `cu102`, or `cu111` depending on your PyTorch installation.
rusty1s's avatar
rusty1s committed
65

66
67
68
69
70
|             | `cpu` | `cu101` | `cu102` | `cu111` |
|-------------|-------|---------|---------|---------|
| **Linux**   | ✅    | ✅      | ✅      | ✅      |
| **Windows** | ✅    | ❌      | ✅      | ✅      |
| **macOS**   | ✅    |         |         |         |
rusty1s's avatar
rusty1s committed
71

72
**Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0 and PyTorch 1.7.0/1.7.1 (following the same procedure).
rusty1s's avatar
rusty1s committed
73
74
75
76

### 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
77
78
79

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

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

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

rusty1s's avatar
bugfix  
rusty1s committed
89
90
91
Then run:

```
rusty1s's avatar
rusty1s committed
92
pip install torch-spline-conv
rusty1s's avatar
bugfix  
rusty1s committed
93
94
```

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

rusty1s's avatar
bugfix  
rusty1s committed
102
103
104
## Usage

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

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

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

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

### Returns

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

### Example

```python
import torch
rusty1s's avatar
rusty1s committed
152
from torch_spline_conv import spline_conv
rusty1s's avatar
bugfix  
rusty1s committed
153

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

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

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

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

## Running tests

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

## 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
```