PyTorch Geometric: A Fast PyTorch Library for DL

A new GitHub project, PyTorch Geometric (PyG), is attracting attention across the machine learning community. PyG is a geometric deep learning extension library for PyTorch dedicated to processing irregularly structured input data such as graphs, point clouds, and manifolds. The project was developed and released by two PhD students from TU Dortmund University, Matthias Fey and Jan E. Lenssen.

Graph Neural Networks (GNNs) have developed into an effective approach for representation learning on graphs, point clouds and manifolds. GNN implementation is however challenging, as it requires GPUs to process a large amount of highly sparse and irregular data of different sizes.

Fey and Lenssen explain that by leveraging dedicated CUDA kernels, PyG can achieve high performance in runtime experiments. They demonstrate a way to simplify the implementation of GNNs, for example by implementing a single layer like the edge convolution layer:

To compare PyG with other deep graph neural net libraries, Fey and Lenssen conducted several experiments to determine the runtime of the training process on a single NVIDIA GTX 1080 Ti, concluding that “PyG is very fast despite working on sparse data. Compared to the Deep GraphLibrary (DGL) 0.1.3, PyG trains models up to 15 times faster.”

The project has received over 2000 stars on GitHub and won widespread praise from the machine learning community. Facebook AI Chief Yann LeCun endorsed PyG as “a fast & nice-looking PyTorch library.” More information about PyG is available on the project’s GitHub page.

Author: Fangyu Cai | Editor: Michael Sarazen