Welcome to repository for Hierarchical Graph Neural Network for Particle Track Reconstruction.
- To present a Hierarchical GNN implementation to the HEP and ML community.
- To present an example of using HGNN for particle tracking on the TrackML dataset.
- To provide a quick comparison between HGNN and flat GNNs.
It's recommended to start a conda environment before installation:
conda create --name exatrkx-tracking python=3.8
conda activate exatrkx-tracking
pip install pip --upgrade
If you have a CUDA GPU available, load the toolkit or install it now. You should check that this is done by running nvcc --version. Then, install the following dependencies:
cudf 22.04.00
cugraph 22.04.00
cuml 22.04.00
cupy 9.6.0
frnn 0.0.0
numba 0.55.1
numpy 1.21.6
pytorch-lightning 1.6.3
numpy 1.21.6
rapids 22.04.00
scikit-learn 1.0.2
scipy 1.8.0
torch 1.11.0+cu115
torch-geometric 2.0.4
torch-scatter 2.0.9
torch-sparse 0.6.13
wandb 0.12.16
yaml 0.2.5
Other combinations of version should also be working but not have tested yet. Finally, clone the repo and you are ready for quick-start.
Once installation is done, download the quickstart TrackML 1GeV filter-processed dataset. Then, change the line input_dir: /global/cfs/cdirs/m3443/data/trackml-codalab/gnn_processed/1GeV_pt_cut_quickstart_example here to the directory where you expand the file. After that, navigate to the example notebook. You should be able to see a variable called ROOT_PATH. Set it to be the directory you wish to keep model checkpoints and loggings. Then run the section import and training a new model. There will be an input box that you should enter model ID or model name. Use 4 to try bipartite classifier HGNN for our first try. After that, a WandB login token will be needed for logging purpose. Provide yours if you already have one or register a new account at wandb. Or alternatively you can comment out logger = WandbLogger(project="TrackML_1GeV") and change it to logger = None to disable logging. The model requires quite a lot GPU memory to run, so if cuda runs out of memory, navigate to configs and change latent till the model fits to your GPU. Some advanced training techniques (e.g. multiple GPU training) can be found here but note that we have observed some issues about multi-gpu training and are still working on to stablize it.