/pytorch_geometric

Geometric Deep Learning Extension Library for PyTorch

Primary LanguagePythonMIT LicenseMIT


PyPI Version Build Status Code Coverage

Documentation | Paper

PyTorch Geometric (PyG) 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, multi gpu-support, 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.


PyTorch Geometric makes implementing Graph Neural Networks a breeze (see here for the accompanying tutorial). For example, this is all it takes to implement the edge convolutional layer:

import torch
from torch.nn import Sequential as Seq, Linear as Lin, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
    def __init__(self, F_in, F_out):
        super(EdgeConv, self).__init__(aggr='max')  # "Max" aggregation.
        self.mlp = Seq(Lin(2 * F_in, F_out), ReLU(), Lin(F_out, F_out))

    def forward(self, x, edge_index):
        # x has shape [N, F_in]
        # edge_index has shape [2, E]
        return self.propagate(edge_index, x=x)  # shape [N, F_out]

    def message(self, x_i, x_j):
        # x_i has shape [E, F_in]
        # x_j has shape [E, F_in]
        edge_features = torch.cat([x_i, x_j - x_i], dim=1)  # shape [E, 2 * F_in]
        return self.mlp(edge_features)  # shape [E, F_out]

In addition, PyTorch Geometric is fast, even compared to other deep graph neural net libraries:

Dataset Epochs Model DGL-DB DGL-SPMV PyG
Cora 200 GCN 4.19s 0.32s 0.25s
GAT 6.31s 5.36s 0.80s
CiteSeer 200 GCN 3.78s 0.34s 0.30s
GAT 5.61s 4.91s 0.88s
PubMed 200 GCN 12.91s 0.36s 0.32s
GAT 18.69s 13.76s 2.42s
MUTAG 200 RGCN 18.81s 2.40s 2.14s
DGL-DB: DGL 0.2 with Degree Bucketing.
DGL-SPMV: DGL 0.2 with SPMV optimization inspired by PyG.
Training runtimes obtained on a NVIDIA GTX 1080Ti.

In detail, the following methods are currently implemented:


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 examples in the examples/ directory.

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.

Installation

Ensure that at least PyTorch 1.0.0 is installed and verify that cuda/bin, cuda/include and cuda/lib64 are in your $PATH, $CPATH and $LD_LIBRARY_PATH 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:...

and

$ echo $LD_LIBRARY_PATH
>>> /usr/local/cuda/lib64

on Linux or

$ echo $DYLD_LIBRARY_PATH
>>> /usr/local/cuda/lib

on macOS. Then run:

$ pip install --verbose --no-cache-dir torch-scatter
$ pip install --verbose --no-cache-dir torch-sparse
$ pip install --verbose --no-cache-dir torch-cluster
$ pip install --verbose --no-cache-dir torch-spline-conv (optional)
$ pip install torch-geometric

See the Frequently Asked Questions and verify that your CUDA is set up correctly if you encounter any installation errors.

Running examples

cd examples
python cora.py

Cite

Please cite our paper (and the respective papers of the methods used) if you use this code in your own work:

@inproceedings{Fey/Lenssen/2019,
  title={Fast Graph Representation Learning with {PyTorch Geometric}},
  author={Fey, Matthias and Lenssen, Jan E.},
  booktitle={ICLR Workshop on Representation Learning on Graphs and Manifolds},
  year={2019},
}

Running tests

python setup.py test