Gabriele Cosmo
108 lines
4.6 KiB
Markdown
108 lines
4.6 KiB
Markdown
\page refPar04training_vae Module training_vae
|
|
|
|
This repository contains the set of scripts used to train, generate and validate the generative model used
|
|
in this example.
|
|
|
|
- root2h5_for_vae.py: translation of ROOT file with showers to h5 files usable in the VAE model.
|
|
- root2h5.py: translation of ROOT file with showers to h5 files, more general version (recommended). It allows to simulate non-discrete energies and stores showers in 3D tensors (R x phi x z).
|
|
- core/constants.py: defines the set of common variables.
|
|
- core/model.py: defines the VAE model class and a handler to construct the model.
|
|
- utils/preprocess.py: defines the data loading and preprocessing functions.
|
|
- utils/hyperparameter_tuner.py: defines the HyperparameterTuner class.
|
|
- utils/gpu_limiter.py: defines a logic responsible for GPU memory management.
|
|
- utils/observables.py: defines a set of observable possibly calculated from a shower.
|
|
- utils/plotter.py: defines plotting classes responsible for manufacturing various plots of observables.
|
|
- train.py: performs model training.
|
|
- generate.py: generate showers using a saved VAE model.
|
|
- observables.py: defines a set of shower observables.
|
|
- validate.py: creates validation plots using shower observables.
|
|
- convert.py: defines the conversion function to an ONNX file.
|
|
- tune_model.py: performs hyperparameters optimization.
|
|
|
|
## Getting Started
|
|
|
|
`setup.py` script creates necessary folders used to save model checkpoints, generate showers and validation plots.
|
|
|
|
```
|
|
python3 setup.py
|
|
```
|
|
|
|
## Full simulation dataset
|
|
|
|
The full simulation dataset can be downloaded from/linked to [Zenodo](https://zenodo.org/record/6082201#.Ypo5UeDRaL4).
|
|
|
|
If custom simulation is used, the output of full simulation must be translated to h5 files using `root2h5_for_vae.py` script. This file is recommended for use with the provided VAE model. For all other usecases script `root2h5.py` is recommended, as it does not assume that simulation is run with discrete energies (e.g. GPS can be used within Geant4 simulation instead of the particle gun).
|
|
|
|
## Training
|
|
|
|
In order to launch the training:
|
|
|
|
```
|
|
python3 train.py
|
|
```
|
|
|
|
You may specify those three following flags. If you do not, then default values will be used.
|
|
|
|
```--max-gpu-memory-allocation``` specifies a maximum memory allocation on a single, logic GPU unit. Should be given as
|
|
an integer.
|
|
|
|
```--gpu-ids``` specifies IDs of physical GPUs. Should be given as a string, separated with comas, no spaces.
|
|
If you specify more than one GPU then automatically ```tf.distribute.MirroredStrategy``` will be applied to the
|
|
training.
|
|
|
|
```--study-name``` specifies a study name. This name is used as an experiment name in W&B dashboard and as a name of
|
|
directory for saving models.
|
|
|
|
See ```run.sh``` and ```condor.sub``` for training on HTCondor.
|
|
Note: wandb api key is hardcoded and needs to be added manually in order to log stats to weights and biases.
|
|
|
|
## Hyperparameters tuning
|
|
|
|
If you want to tune hyperparameters, specify in `tune_model.py` parameters to be tuned. There are three types of
|
|
parameters: discrete, continuous and categorical. Discrete and continuous require range specification (low, high), while
|
|
the categorical parameter requires a list of possible values to be chosen. Then run it with:
|
|
|
|
```
|
|
python3 tune_model.py
|
|
```
|
|
|
|
If you want to parallelize tuning process you need to specify a common storage (preferable MySQL database) by
|
|
setting `--storage="URL_TO_MYSQL_DATABASE"`. Then you can run multiple processes with the same command:
|
|
|
|
```
|
|
python3 tune_model.py --storage="URL_TO_MYSQL_DATABASE"
|
|
```
|
|
|
|
Similarly to training procedure, you may specify ```--max-gpu-memory-allocation```, ```--gpu-ids``` and
|
|
```--study-name```.
|
|
|
|
## ML shower generation (MLFastSim)
|
|
|
|
In order to generate showers using the ML model, use `generate.py` script and specify information of geometry, energy
|
|
and angle of the particle and the epoch of the saved checkpoint model. The number of events to generate can also be
|
|
specified (by default is set to 10.000):
|
|
|
|
```
|
|
python3 generate.py --geometry=SiW --energy=64 --angle=90 --epoch=1000 --study-name=YOUR_STUDY_NAME
|
|
```
|
|
|
|
If you do not specify an epoch number the based model (saved as ```VAEbest```) will be used for shower generation.
|
|
|
|
## Validation
|
|
|
|
In order to validate the MLFastSim and the full simulation, use `validate.py` script and specify information of
|
|
geometry, energy and angle of the particle:
|
|
|
|
```
|
|
python3 validate.py --geometry=SiW --energye=64 --angle=90
|
|
```
|
|
|
|
## Conversion
|
|
|
|
After training and validation, the model can be converted into a format that can be used in C++, such as ONNX,
|
|
use `convert.py` script:
|
|
|
|
```
|
|
python3 convert.py --epoch 1000
|
|
```
|