Accelerating the Inference Time of Machine Learning-based Track Finding Pipeline
Date:
Recently, graph neural networks (GNNs) have been successfully used for a variety of reconstruction problems in HEP. In this work, we develop and evaluate an end-to-end C++ implementation for inferencing a charged particle tracking pipeline based on GNNs. The pipeline steps include data encoding, graph building, edge filtering, GNN and track labeling and it runs on both GPUs and CPUs. The ONNX Runtime C++ API is used to run PyTorch deep learning models converted to ONNX. The implementation features an improved GPU-based fixed radius nearest neighbor search for identifying edges and a weakly connected component algorithm for the labeling step. In addition, complete conversion to C++ allows integration with existing tracking software, including ACTS. We report the memory usage, average event latency, and the efficiency and purity tracking performance of our implementation applied to the TrackML benchmark dataset. The GPU-based implementation provides considerable speed-ups over the CPU-based execution and can be extended to run on multiple GPUs.