Stop Thinking, Just Do!

Sung-Soo Kim's Blog

PyTorch Geometric


17 February 2019

Article Source

PyPI Version Build Status Code Coverage


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.

PyTorch Geometric makes implementing graph convolutional networks a breeze (see here for the accompanying tutorial). For example, this is all it takes to implement a single layer like the edge convolution 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__()
        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(aggr='max', edge_index=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 =[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 PyTorch Geometric
Cora 200 GCN 4.2s 0.7s
GAT 33.4s 2.2s
CiteSeer 200 GCN 3.9s 0.8s
GAT 28.9s 2.4s
PubMed 200 GCN 12.7s 2.0s
GAT 87.7s 12.3s
MUTAG 50 R-GCN 3.3s 2.4s
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 provided 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.


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 (optional)
$ 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.

Docker image

You can also run PyTorch Geometric with CUDA-9.0 inside a docker image:

$ docker pull shengwenliang/pytorch_graph

Running examples

cd examples

Running tests

python test

comments powered by Disqus