Cargando…

A ROOT feature for parsing PyTorch Geometric graph neural networks into C++ code for fast inference

Graph neural networks have proven effective in many different fields — including particle physics — and are typically trained using Python tools such as PyTorch Geometric. For deployment, however, it is often desirable to move away from Python and run fast inference in a high-performance environment...

Descripción completa

Detalles Bibliográficos
Autor principal: Van Berkum, Stefan
Lenguaje:eng
Publicado: 2023
Materias:
Acceso en línea:http://cds.cern.ch/record/2868483
Descripción
Sumario:Graph neural networks have proven effective in many different fields — including particle physics — and are typically trained using Python tools such as PyTorch Geometric. For deployment, however, it is often desirable to move away from Python and run fast inference in a high-performance environment like C++. This report describes my work on adding a ROOT feature for parsing PyTorch Geometric graph neural networks into C++ code. The feature currently supports basic graph neural networks and preliminary results suggest that it is roughly five times faster than TorchScript (PyTorch's native C++ inference code generator) on a single CPU thread, for a basic graph-classification model with two graph convolutions. In the future, additional modules should be added to support the entire range of GNNs used in LHC experiments and further optimizations could yield even larger speedups.