TY - JOUR
T1 - Accelerating the Inference of the Exa.TrkX Pipeline
AU - Lazar, Alina
AU - Ju, Xiangyang
AU - Murnane, Daniel
AU - Calafiura, Paolo
AU - Farrell, Steven
AU - Xu, Yaoyuan
AU - Spiropulu, Maria
AU - Vlimant, Jean Roch
AU - Cerati, Giuseppe
AU - Gray, Lindsey
AU - Klijnsma, Thomas
AU - Kowalkowski, Jim
AU - Atkinson, Markus
AU - Neubauer, Mark
AU - DeZoort, Gage
AU - Thais, Savannah
AU - Hsu, Shih Chieh
AU - Aurisano, Adam
AU - Hewes, Jeremy
AU - Ballow, Alexandra
AU - Acharya, Nirajan
AU - Wang, Chun Yi
AU - Liu, Emma
AU - Lucas, Alberto
N1 - Publisher Copyright:
© Published under licence by IOP Publishing Ltd.
PY - 2023
Y1 - 2023
N2 - Recently, graph neural networks (GNNs) have been successfully used for a variety of particle reconstruction problems in high energy physics, including particle tracking. The Exa.TrkX pipeline based on GNNs demonstrated promising performance in reconstructing particle tracks in dense environments. It includes five discrete steps: data encoding, graph building, edge filtering, GNN, and track labeling. All steps were written in Python and run on both GPUs and CPUs. In this work, we accelerate the Python implementation of the pipeline through customized and commercial GPU-enabled software libraries, and develop a C++ implementation for inferencing the pipeline. The implementation features an improved, CUDA-enabled fixed-radius nearest neighbor search for graph building and a weakly connected component graph algorithm for track labeling. GNNs and other trained deep learning models are converted to ONNX and inferenced via the ONNX Runtime C++ API. The complete C++ implementation of the pipeline allows integration with existing tracking software. We report the memory usage and average event latency tracking performance of our implementation applied to the TrackML benchmark dataset.
AB - Recently, graph neural networks (GNNs) have been successfully used for a variety of particle reconstruction problems in high energy physics, including particle tracking. The Exa.TrkX pipeline based on GNNs demonstrated promising performance in reconstructing particle tracks in dense environments. It includes five discrete steps: data encoding, graph building, edge filtering, GNN, and track labeling. All steps were written in Python and run on both GPUs and CPUs. In this work, we accelerate the Python implementation of the pipeline through customized and commercial GPU-enabled software libraries, and develop a C++ implementation for inferencing the pipeline. The implementation features an improved, CUDA-enabled fixed-radius nearest neighbor search for graph building and a weakly connected component graph algorithm for track labeling. GNNs and other trained deep learning models are converted to ONNX and inferenced via the ONNX Runtime C++ API. The complete C++ implementation of the pipeline allows integration with existing tracking software. We report the memory usage and average event latency tracking performance of our implementation applied to the TrackML benchmark dataset.
UR - http://www.scopus.com/inward/record.url?scp=85149773021&partnerID=8YFLogxK
UR - http://www.scopus.com/inward/citedby.url?scp=85149773021&partnerID=8YFLogxK
U2 - 10.1088/1742-6596/2438/1/012008
DO - 10.1088/1742-6596/2438/1/012008
M3 - Conference article
AN - SCOPUS:85149773021
SN - 1742-6588
VL - 2438
JO - Journal of Physics: Conference Series
JF - Journal of Physics: Conference Series
IS - 1
M1 - 012008
T2 - 20th International Workshop on Advanced Computing and Analysis Techniques in Physics Research, ACAT 2021
Y2 - 29 November 2021 through 3 December 2021
ER -