Oct 19 – 23, 2020
Europe/Zurich timezone

Graph Convolutional Operators in the the PyTorch JIT

Oct 23, 2020, 3:10 PM
Lightning talk 6 ML infrastructure : Hardware and software for Machine Learning Workshop


Lindsey Gray (Fermi National Accelerator Lab. (US))


The PyTorch just-in-time (jit) compiler is a powerful tool for optimizing and serializing neural network models. However, its range is limited by the subset of the python language that it is restricted to and the number of tensor operations implemented in C++. These limitations were a major blocker to using graph neural networks implemented in the geometric deep learning (GDL) library PyTorch Geometric (PyG) at scale. In particular, models being researched needed to be re-implemented, validated, and retrained before they could be deployed in inference as a service (IaaS) frameworks for wider use. To solve this, the PyG framework was extended to include an automatic analysis of user-defined convolutional operators that renders structurally identical, concrete, and jit-compatible versions of the operator. The results of these additions are GDL models that can seamlessly flow from research-focused workflows to development and deployment at scale with standard IaaS infrastructure. In this presentation we will discuss the additions made to PyG to achieve this new functionality and give examples with preliminary performance estimates of live models in the context of high energy particle physics.

Primary author

Lindsey Gray (Fermi National Accelerator Lab. (US))


Thomas Klijnsma (Fermi National Accelerator Lab. (US)) Matthias Fey (TU Dortmund)

Presentation materials