pip install omnilearned

A few standard datasets can be directly downloaded using the command:

omnilearned dataloader -d DATASET -f OUTPUT/PATH

Datasets available are: top/qg/aspen/atlas/jetclass/h1 If --d pretrain is used instead, aspen, atlas, jetclass, and h1 datasets will be downloaded. The total size of the pretrain dataset is around 4T so be sure to have enough space available.

Single GPU training can be started using:

omnilearned train  -o ./ --save_tag test --dataset DATASET --path OUTPUT/PATH

For multiple GPUs and SLURM you can use the train.sh example script

#Inside an interactive session run
./train.sh

To lint the code, run:

black .
ruff check --fix .