WP21 Train

Current development version: 0.10.0. Legacy APIs remain available.

GitLab source repository · PyPI package · Documentation · Report an issue

WP21 Train orchestrates the training and evaluation operations used by the WP21 group. A trainer adapts a model backend; TrainingRunner runs it once or searches its parameters with Optuna. The package also provides backend-independent classification and trigger evaluation tools.

Changes in this version

The 0.10.0 orchestration interface now keeps older project/physics utilities in legacy_support, while preserving their old imports. Model list scans use length_range (legacy length inputs remain readable). ML Utils models can load separate HGQ2 configuration, including scannable quantizer/beta settings. Keras objective adapters preserve optional sample weights from (inputs, targets, sample_weight) batches.

Optuna model-build failures are recorded and skipped up to a configurable limit; see Optuna settings and defaults before starting a study. The custom layer/model hook registers project-owned Python classes before building models. Configure it with meta.custom_objects_factory: custom_layers.py:get_custom_objects.

Installation

Python 3.10 or newer is required. Start with an isolated environment:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

Install the latest published package from PyPI:

python -m pip install wp21_train

Choose the extras for your workflow. These are alternatives; you do not need to install every backend. Omit optuna if you only need standalone training.

For a CPU-only Torch installation, install Torch from its CPU wheel index first. The tensorflow extra includes the historical tf2onnx dependency for legacy export; native .keras saves do not need ONNX. Use separate environments for legacy HGQ and modern HGQ2 when their dependencies conflict.

Check the version installed in your active environment:

python -c 'from importlib.metadata import version; print(version("wp21_train"))'

Core imports, configuration, event contracts and array evaluation do not import TensorFlow, Keras, Torch, HGQ, ML Utils or Optuna. Importing a concrete trainer without its backend raises an installation hint. The base installation includes numerical/data dependencies. Fixed runs do not require Optuna.

Install the development version or contribute

The interfaces documented here target development version 0.10.0 on train_orchestrator_release. Installing from PyPI selects the published release, which may differ. To use this development branch, clone the GitLab repository and install it in editable mode with your chosen extras:

git clone --branch train_orchestrator_release https://gitlab.cern.ch/atlas-nextgen-wp21/wp21_training_framework.git
cd wp21_training_framework
python -m pip install -e '.[tensorflow,optuna,dev]'

For example, replace tensorflow with brevitas for quantized Torch training. GitLab access may require your CERN account. The documentation links below point to this development branch. Cloning also provides the examples, which are not installed as part of the Python package. Record the commit and dependency versions when reproducing an experiment.

Modular use and compatibility

The components can be used separately. Call an objective on arrays, call a trainer directly, or use TrainingRunner when search orchestration is useful. There is no mandatory centralized Project class or runner around standalone training/evaluation. Each adapter still requires its documented model/data input.

Version 0.10.0 adds wp21_train.objectives.Objective, a backend-independent base for user-defined metrics, plus generic MultiClassAUC and TriggerEfficiency. Prediction bindings and optional Keras/ML Utils epoch adapters reuse the same calculations. See generic objectives and complete examples.

The current dummy jet-tagger's imports, objective factory, params.yml and model.yml remain supported without migration. Its legacy MultiClassAUC and ObjectiveTuner path is retained. Updating the package does not switch the framework to the new runner. The training image must install the new package version to use new APIs.

Supported interfaces

All five implement Trainer: prepare, fit, train, predict, and save. Factories must construct fresh models. A saved Keras model is recompiled with a fresh optimizer; Torch checkpoint loading requires the matching factory. ML Utils handles its own supported quantization configuration, including use_hgq. There is no automatic dispatch between these adapters yet.

Legacy TrainerBase, LossTuner, ObjectiveTuner, HGQ 0.2, XGBoost, export and physics utilities remain available. They are not all migrated to the new interface. HGQ2 and legacy HGQ use different APIs/environments.

Input and output reference

See input and output types for array shapes, data types, units, model factories and returned objects. Class docstrings describe constructor inputs, defaults, shapes and lifecycle behavior, including legacy utilities. Inspect them directly from Python:

from wp21_train.training import TrainingRunner, TrainingResult
from wp21_train.evaluation import turn_on_curve

help(TrainingRunner)
help(TrainingResult)
help(turn_on_curve)

Configuration and training

The existing YAML/JSON data/meta structure is preserved. For native trainers, data contains parameter choices. Save this example as params.yml:

data:
  lr: [0.001, 0.0001]
  batch_size: [64]
  epochs: [10]
  optimiser: [adam]
  loss: [sparse_categorical_crossentropy]
  metrics: [accuracy]
meta:
  model_config: model.yml
  search: optuna
  sampler: tpe
  trials: 4
  objective: val_loss
  direction: minimize
  seed: 42
  callbacks:
    early_stopping:
      monitor: val_loss
      mode: minimize
      patience: 2
      min_delta: 0.0
      restore_best_weights: true

MLUtilsTrainer instead reads architecture and dynamic scan specifications from meta.model_config. Other trainers currently receive their model source explicitly in Python; model_config does not route them automatically. Paths supplied to a trainer resolve relative to its configuration file. No new Project subclass is required.

from wp21_train.training import TrainingRunner
from wp21_train.training.trainers import KerasTrainer

# Supply dataset with x_train, y_train, x_val and y_val attributes.
# Supply model_factory returning a fresh compiled-or-uncompiled Keras model.
# A runnable example is provided in examples/mnist/run.py (see below).
trainer = KerasTrainer('params.yml', dataset, model_factory)
runner  = TrainingRunner(trainer)
results = runner.run()
winner  = runner.best_result
if winner is not None:
    trainer.save('output/winner', winner)

sampler accepts tpe/bayesian, random, or grid. Python callers can inject an Optuna BaseSampler instance. search: grid and search: random retain legacy flat-scan ordering and no-replacement semantics; explicit Optuna random sampling can repeat choices. Callable model-defined grid searches need finite meta.grid_space choices. See search orchestration.

To train once, pass concrete values to trainer.train(parameters) or runner.run(parameters). Missing/fixed/none search settings require singleton choices. A fixed ML Utils model automatically bypasses search. No optimization objective is required for fixed runs.

Results retain the model, parameters, configuration, history, scalar metrics and trial context. Searches use the metric named by meta.objective or an explicit objective(result) callable. The winner is an already trained result; saving it does not retrain it. Retaining every trial model can consume substantial memory.

Brevitas and Optuna

BrevitasTrainer reuses the Torch lifecycle with a factory receiving concrete Optuna parameters and fixed quantizer settings. Weight/activation bit widths can be scanned alongside ordinary training parameters. It supports standalone QAT, generic objectives, early stopping, pruning and native checkpoint reloads. See Brevitas factory/configuration examples. PTQ, quantized ONNX/QONNX export and framework synthesis integration are separate capabilities.

Callbacks and evaluation

Common training supports fresh per-trial event callbacks and explicit early stopping across all five adapters. Native adapters apply these policies through the common train/fit(context) interface. Legacy calls keep their historical callback behavior. The ML Utils legacy callback factory remains supported.

from wp21_train.callbacks import TrainingCallback

class Progress(TrainingCallback):
    def on_epoch_end(self, event):
        print(event.context.trial_number, event.epoch, dict(event.metrics))

trainer = KerasTrainer('params.yml', dataset, model_factory,
                      event_callbacks_factory=lambda context: [Progress()])

Restored weights and result metrics refer to the same epoch. Only model weights are restored, not optimizer state. Optional Optuna pruning hooks report epoch metrics to a trial. See callbacks and evaluation for pruning, event ordering and supported policy options.

Evaluation consumes explicit score arrays; it does not assume a jet tagger, a class index or a probability conversion:

from wp21_train.evaluation import binary_auc, rate_curve, turn_on_curve

auc   = binary_auc(binary_labels, positive_class_scores)
rates = rate_curve(background_scores, thresholds, input_rate_hz=40_000_000)
turnon = turn_on_curve(reference_pt, positive_class_scores, threshold=0.5,
                      bins=[0, 20, 40, 60, 100])
rates.save('output/rate.json')
turnon.save('output/turnon.json')

Rate normalization is explicit in Hz. Selections use score >= threshold. Nonnegative event weights are supported; empty turn-on bins serialize as null. BinaryAUCObjective connects prediction to scalar search scoring, with explicit output selection for named/multiclass outputs. trigger_rate returns a scalar rate at a chosen threshold. Curves are separate artifacts, not scalar objectives. There are no inferred physics selections or weighted-error prescriptions.

Examples

MNIST examples cover all five adapters. From the repository root after cloning and installing the corresponding extras above:

python examples/mnist/run.py --backend keras --download --orchestrated
python examples/mnist/run.py --backend torch --data mnist.npz --standalone

Downloads happen only with --download. The default example path exercises legacy APIs; --standalone uses the common trainer and --orchestrated uses TrainingRunner.

Scope and documentation

Automatic model routing, framework/MLflow pipeline migration, distributed study resume and multi-objective search are not implemented by this interface yet. Native checkpoints support model reuse/inference, not exact optimizer resume. Save project code, input data and dependency versions alongside model artifacts when reproducibility is required.

See the documentation index for API contracts, input/output types and compatibility details. The historical API reference retains the previous README's legacy examples and utilities.

Developers and contributors

The developers credited in the previous README are:

• Ioannis Xiotidis — ioannis.xiotidis@cern.ch
• Pawel Mucha — pawel.mucha@cern.ch
• Vila Andela Petrovic — vila.andela.petrovic@cern.ch
• David Reikher — david.reikher@cern.ch