Diffusion4Sim (CaloDiT)

Description

With this repository you can train and evaluate diffusion models for fast simulation of calorimeter showers. Currently the following methods are supported:

• DDPM^{1 2} and DDIM³
• EDM⁴
• DPM-Solver⁵ and DPM-Solver++⁶
• Consistency Distillation⁷
• Easy Consistency Tuning⁸

It also features:

• DiT⁹-like transformer architecture
• multi-GPU training with Accelerate
• training progress tracking with Weights & Biases
• conversion to c++ inference compatible forms (ONNX & TorchScript)
• training, distillation, validation and generation pipelines
• ...and more

To train the models, you can use one of the datasets provide in the CaloChallenge. Keep in mind that some changes in the src/data/geometry.py file might be needed to load the data properly.

Note: Model available in Par04 is based on EDM and distilled using consistency distillation.

Installation

# clone repository
git clone https://gitlab.cern.ch/fastsim/diffusion4sim.git
cd diffusion4sim

# create environment and install dependencies
conda create -n diffusion4sim python=3.9
conda activate diffusion4sim
pip install -r requirements.txt

# copy the examplar .env file and fill in the paths and W&B credentials
cp .env-example .env && vim .env

Usage

Training

Modify the appropriate configuration file in the configs/train directory to specify all hyperparameters. Refer configs/train/edm_allegro_scratch.yaml Then, run the training script with the path to the config file as the first argument. You can also specify command line arguments to override the config values. For example:

python scripts/train.py configs/train/edm.yaml experiment.run_name="edm_calo_challenge_dataset2"

The training script will save the model checkpoints, logs and intermediate evaluation results in the specified output_dir directory. You can track the training progress with Weights & Biases by setting the use_wandb parameter in the config file.

To run multi-gpu training, use the following command:

accelerate launch --multi_gpu --num_processes=<num_gpus> scripts/train.py configs/train/edm_allegro_scratch.yaml

For more information, see the Accelerate documentation.

Distillation

Similarly to training, you run distillation using the scripts/distill.py script. The distillation process requires the teacher model checkpoint, which should be specified in the config file as model_path. Currently only the Consistency Distillation⁷ is supported. Refer configs/distill/cd_allegro_scratch.yaml.

Conversion

Conversion to ONNX or TorchScript models can be done using scripts/convert.py. Refer the file as an example.

Validation

Use the scripts/validate.py to generate showers under the same conditions as in the provided file, and compare with this reference data. The script will save generated showers and create the plots comparing observables of generated samples against reference showers. See the run_evaluation.sh script for an example.

Generation

Use the scripts/generate.py to generate showers from the trained model under the specified conditions. The generated showers will be saved in the output_dir directory in the .h5 format for further analysis.

Multi-geometry training and adaptation

Training

If you need to pretrain on multiple geometries, follow the instructions below. A couple of changes to inform about the geometry is needed, mostly in the .yaml files. See configs/train/edm_multi.yaml for an example.

• Change the way data files are described both during training and testing. Use [geometry_name, file_path] instead of file_path.
• Setup need_geo_condn & train_on. In model architecture, change conditions_size. This should be one more than the length of train_on if you are interested in finetuning over a new detector later.
• Recommended to use the standarization based on all of your pretraining geometries.

Distillation: Make the above changes for the distillation .yaml file and mention the teacher checkpoint of the pretrained model as stated previously. Refer configs/distill/cd_multi.yaml

Adaptation

You would need a pretrained model to finetune. Get one from this link or following the previous steps. Adaptation is done by re-training the pretrained model using a low learning rate. So .yaml file will be same as the one used in multi-geometry training, except:

• Don't use a very high learning rate. Anything less than or equal to 1e-3 should be fine.
• Fewer training steps, i.e., max_steps and respective lr_schedular args.
• Data files. Keep train_on the same as before. Any detector geometry not in the list will be the last one, i.e., [0 0 ... 1].
• Mention the checkpoint on the pretrained model as model_path under model params.
• Refer to configs/train/edm_allegro_adapt.yaml

Distillation:

• If you are not doing the pretraining yourselves, get the pretrained distilled model from this link.
• Mention checkpoints in .yaml file for the distillation. Teacher checkpoint is pretrained EDM model. Student checkpoint is pretrained distilled model.
• Make sure init_student_from_teacher is false.
• Refer configs/distill/cd_allegro_adapt.yaml for more details.

Credits

This repository is derived from the Diffusion4FastSim, however it also benefits a lot from the code implemented in the following repositories:

• denoising-diffusion-pytorch
• improved-diffusion
• k-diffusion
• consistency_models
• dpm-solver
• DiT