Welcome to code repository for Efficient and Robust Jet Tagging at the LHC with Knowledge Distillation.
git clone https://github.com/ryanliu30/KD4Jets.git --recurse-submodules
cd KD4Jets
conda create -n KD4Jets python=3.11
conda activate KD4Jets
pip install -r requirements.txt
pip install -e .
Then, please follow the instructions in LorentzNet to get the training data from OSF and put them under the data
directory.
To begin with, run the following command:
python train.py --cfg experiments/deepset.yaml
This will train a deepset jet tagging model from scratch. To train with knowledge distillation, run:
python train.py --cfg experiments/deepsetKD.yaml
Note that the only difference of these two config file is that MLP.yaml
and MLPKD.yaml
that can be run using the same command.
To evaluate the models, run
python eval.py --cfg plots/example.yaml
Where in the yaml
file you should specify the checkpoint to use. They should locate under the LorentzNetKD
folder. We also supplied four model checkpoints from our training that can be used out of the box.
If you use this work in your research, please cite:
@misc{liu2023efficient,
title={Efficient and Robust Jet Tagging at the LHC with Knowledge Distillation},
author={Ryan Liu and Abhijith Gandrakota and Jennifer Ngadiuba and Maria Spiropulu and Jean-Roch Vlimant},
year={2023},
eprint={2311.14160},
archivePrefix={arXiv},
primaryClass={hep-ex}
}