Speakers
Description
Graph Neural Network (GNN) models proved to perform well on the particle track finding problem, where traditional algorithms become computationally complex as the number of particles increases, limiting the overall performance. GNNs can capture complex relationships in event data represented as graphs. However, training on large graphs is challenging due to computation and GPU memory requirements. The graph representation must fit into the GPU memory to fully utilize event data when training the GNN model. Otherwise, the graphs must be divided into smaller batches. We evaluate generic sampling methods that modify the conventional GNN training by using the mini-batch scheme to reduce the amount of required memory and facilitate parallel processing. Although splitting graphs may seem straightforward, striking a balance between computational efficiency and preserving the essential characteristics of the graph is challenging.
Through empirical experiments, we aim to test and tune graph sampling and partitioning methods to improve the edge classification performance for track reconstruction. Node, edge, and subgraph sampling methods are explored to divide data into smaller mini-batches for training. Preliminary results on the TrackML dataset show performance similar to full-batch training. These results prove the effectiveness of sampling methods in edge-level GNN classification tasks and the possibility of extending training to event graphs exceeding the top-of-the-line GPU’s memory for improved performance.