PyTorch Geometric is a geometric deep learning extension library for PyTorch.
It consists of various methods for deep learning on graphs and other irregular structures, also known as geometric deep learning, from a variety of published papers. In addition, it consists of an easy-to-use mini-batch loader, a large number of common benchmark datasets (based on simple interfaces to create your own), and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.
In detail, the following methods are currently implemented:
- SplineConv from Fey et al.: SplineCNN: Fast Geometric Deep Learning with Continuous B-Spline Kernels (CVPR 2018)
- GCNConv from Kipf and Welling: Semi-Supervised Classification with Graph Convolutional Networks (ICLR 2017)
- ChebConv from Defferrard et al.: Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering (NIPS 2016)
- NNConv adapted from Gilmer et al.: Neural Message Passing for Quantum Chemistry (ICML 2017)
- GATConv from Veličković et al.: Graph Attention Networks (ICLR 2018)
- AGNNConv from Thekumparampil et al.: Attention-based Graph Neural Network for Semi-Supervised Learning (CoRR 2017)
- SAGEConv from Hamilton et al.: Inductive Representation Learning on Large Graphs (NIPS 2017)
- GraphConv from, e.g., Morris et al.: Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks (AAAI 2019)
- GINConv from Xu et al.: How Powerful are Graph Neural Networks? (ICLR 2019, under review)
- EdgeConv from Wang et al.: Dynamic Graph CNN for Learning on Point Clouds (CoRR, 2018)
- PointConv from Qi et al.: PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation (CVPR 2017) and PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space (NIPS 2017)
- XConv from Li et al.: PointCNN: Convolution On X-Transformed Points (NeurIPS 2018)
- GMMConv from Monti et al.: Geometric Deep Learning on Graphs and Manifolds using Mixture Model CNNs (CVPR 2017) and PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space (CVPR 2017)
- A MetaLayer for building any kind of graph network similar to the TensorFlow Graph Nets library from Battaglia et al.: Relational Inductive Biases, Deep Learning, and Graph Networks (CoRR 2018)
- Set2Set from Vinyals et al.: Order Matters: Sequence to Sequence for Sets (ICLR 2016)
- Sort Pool from Zhang et al.: An End-to-End Deep Learning Architecture for Graph Classification (AAAI 2018)
- Dense Differentiable Pooling from Ying et al.: Hierarchical Graph Representation Learning with Differentiable Pooling (NeurIPS 2018)
- Graclus Pooling from Dhillon et al.: Weighted Graph Cuts without Eigenvectors: A Multilevel Approach (PAMI 2007)
- Voxel Grid Pooling from, e.g., Simonovsky and Komodakis: Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs (CVPR 2017)
- Top-K Pooling from Anonymous: Graph U-Net (ICLR 2019, under review) and Cangea et al.: Towards Sparse Hierarchical Graph Classifiers (NeurIPS-W 2018)
- Example of Deep Graph Infomax on Cora from Veličković et al.: Deep Graph Infomax (ICLR 2019, under review)
Head over to our documentation to find out more about installation, data handling, creation of datasets and a full list of implemented methods, transforms, and datasets.
For a quick start, check out our provided examples in the examples/
directory.
We are currently in our first alpha release and work on completing documentation. If you notice anything unexpected, please open an issue and let us know. If you are missing a specific method, feel free to open a feature request. We are constantly encouraged to make PyTorch Geometric even better.
Ensure that at least PyTorch 1.0.0 is installed and verify that cuda/bin
and cuda/include
are in your $PATH
and $CPATH
respectively, e.g.:
$ python -c "import torch; print(torch.__version__)"
>>> 1.0.0
$ echo $PATH
>>> /usr/local/cuda/bin:...
$ echo $CPATH
>>> /usr/local/cuda/include:...
Then run:
$ pip install --upgrade torch-scatter
$ pip install --upgrade torch-sparse
$ pip install --upgrade torch-cluster
$ pip install --upgrade torch-spline-conv
$ pip install torch-geometric
If you are running into any installation problems, please create an issue. Beforehand, please check that the official extension example runs on your machine.
cd examples
python cora.py
Please cite our paper (and the respective papers of the methods used) 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},
}
python setup.py test