wp21_train

A modular, extensible Python framework designed for managing data parsing, serialization, and metadata tracking in machine learning workflows — especially for hardware-aware applications such as HLS and AIE profiling. Built with physicists and hardware engineers in mind, wp21_train integrates support for CERN ROOT I/O, common formats like JSON and Pickle, and parsing of Xilinx toolchain outputs.

🚀 Features

• ✅ Unified interface for reading/writing training data and metadata
• ✅ Supports JSON, Pickle, YAML and ROOT formats
• ✅ Parsers for:
  • HLS reports (Vivado HLS)
  • AIE profiling reports (Vitis AI Engine)
  • ATHENA configuration
• ✅ A training callback interface to log key events and outputs
• ✅ Lightweight type-to-symbol conversion utility
• ✅ Optuna-based hyperparameter optimization
• ✅ ML Utils model training integration
• ✅ Quantization Aware Training with HGQ2 and model checkpoint selection
• ✅ Versioned with easy integration (__version__)

📦 Included Modules

+-----------------------------------------------+--------------------------------------------------------------+-----------------------------------------+
| Module                                        | Description                                                  | Notes / Deps                            |
+-----------------------------------------------+--------------------------------------------------------------+-----------------------------------------+
| wp21_train/savers/json_adapter.py             | JSON-based serialization                                     | stdlib json                             |
| wp21_train/savers/pickle_adapter.py           | Pickle-based serialization                                   | stdlib pickle                           |
| wp21_train/savers/root_adapter.py             | ROOT I/O serialization                                       | uproot, awkward (no ROOT needed)        |
| wp21_train/savers/yml_adapter.py              | YAML I/O serialization                                       | PyYAML (yaml)                           |
| wp21_train/parser/hls_parser.py               | XML parsing of HLS synthesis reports                         |                                         |
| wp21_train/parser/aie_parser.py               | XML parsing of AIE runtime profiling                         |                                         |
| wp21_train/parser/athena_parser.py            | Parsing of ATHENA configuration                              |                                         |
| wp21_train/callbacks/base_callback.py         | Base callback for training pipelines                         |                                         |
| wp21_train/training/searchers/base_search.py  | Base class for search strategies                             |                                         |
| wp21_train/training/searchers/grid_search.py  | Grid search implementation                                   |                                         |
| wp21_train/training/searchers/random_search.py| Random search implementation                                 |                                         |
| wp21_train/training/trainers/base.py          | Common trainer utilities                                     |                                         |
| wp21_train/training/trainers/keras_base.py    | Keras trainer base                                           | tensorflow / keras                      |
| wp21_train/training/trainers/keras_trainer.py | Concrete Keras trainer                                       | tensorflow / keras                      |
| wp21_train/training/trainers/torch_base.py    | PyTorch trainer base                                         | torch                                   |
| wp21_train/training/trainers/torch_trainer.py | Concrete PyTorch trainer                                     | torch                                   |
| wp21_train/training/trainers/hgq_trainer.py   | Concrete HGQ trainer                                         | HGQ                                     |
| wp21_train/training/trainers/hgq2_trainer.py  | Concrete hgq2 trainer                                        | hgq2                                    |
| wp21_train/training/trainers/xgb_trainer.py   | Concrete XGBoost trainer                                     | xgboost                                 |
| wp21_train/training/evaluator.py              | Generic evaluation helpers                                   |                                         |
| wp21_train/training/tuner.py                  | High-level tuning orchestration                              |                                         |
| wp21_train/utils/utility.py                   | Type-shortening / helpers for metadata tagging               |                                         |
| wp21_train/utils/logger.py                    | Simple logging of info/warn/error                            | (renamed from logging.py)               |
| wp21_train/utils/version.py                   | Package versioning (__version__)                             |                                         |
| wp21_train/__init__.py                        | Package exports                                              |                                         |
| wp21_train/training/searchers/optuna_search.py   | Optuna-based hyperparameter search                        | optuna                                  |
| wp21_train/training/trainers/ml_utils_trainer.py | ML Utils trainer backend                                  | wp21_ml_utils                           |
| wp21_train/selection/selector.py                 | Final model/checkpoint selection                          |                                         |
| wp21_train/selection/finalizer.py                | Finalizes and saves the selected ML Utils model           | wp21_ml_utils                           | 
| wp21_train/callbacks/training_callbacks.py       | Callbacks factory for QAT (EBOPs, Parto, BetaScheduler)   | HGQ, TensorFlow                         |
+-----------------------------------------------+--------------------------------------------------------------+-----------------------------------------+

🔧 Installation

From PyPI

pip install wp21_train

📁 Example Usage

🔄 JSON / Pickle / YAML / ROOT Adapters

from wp21_train.savers import json_adapter, pickle_adapter, yml_adapter, root_adapter

adapter = json_adapter("results", dump_data=my_data, dump_meta=my_metadata)
adapter.write_data()

meta, data = adapter.read_data()

🧠 HLS Parser

from wp21_train.parsers import hls_parser

parser = hls_parser("hls_report.xml")
print(parser._data)        # extracted info
print(parser._meta_data)   # associated metadata

⚙️ AIE Parser

from wp21_train.parsers import aie_parser

parser = aie_parser("aie_profile.xml")
print(parser._data)

⚙️ ATHENA Parser

from wp21_train.parsers import athena_parser

parser = athena_parser(data=data_from_adapter, metadata=meta_from_adapter, nevents=10000)
print(parser.config)
print(parser.environment)

📋 Training Callback

from wp21_train.callbacks import base_callback

cb = base_callback(project_name="FastML4Jets")
cb.on_train_begin()
# training loop here
cb.on_train_end()

🌲 XGBoost Trainer

from xgboost import XGBClassifier
import wp21_train as train
from wp21_train.training.trainers import XGBTrainer

def xgb_model(**params):
    params.setdefault("objective", "binary:logistic")
    params.setdefault("eval_metric", ["logloss", "auc"])
    return XGBClassifier(**params)

trainer = XGBTrainer("training_config", dataset, xgb_model)
tuner = train.training.LossTuner("hat", trainer)
tuner.run()

evaluator = train.training.Evaluator(
    tuner.trials,
    trainer,
    meta_path="meta",
    out_path="out",
)
evaluator.save_info(verbose=3)

The XGBoost trainer assumes a scikit-learn-like estimator, such as XGBClassifier or XGBRegressor, with fit(...), evals_result(), and save_model(...)/get_booster().save_model(...) methods. It stores all metrics returned by evals_result() using their native names, such as logloss, auc, rmse, and their val_ counterparts. For LossTuner, only XGBoost logloss is exposed as loss/val_loss. To control which additional metrics are written to history.json, set history_metrics in the config meta section:

{
  "meta": {
    "search": "random",
    "trials": 6,
    "history_metrics": ["auc"]
  }
}

🔍 ML Utils and Optuna Training

wp21_train supports training and hyperparameter optimization of models defined with wp21_ml_utils through MLUtilsTrainer.

The ML Utils workflow uses two YAML configuration files:

• params.yml contains pipeline-level settings such as dataset paths, search strategy, number of trials, optimization direction, and optional selection cuts.
• model.yml contains the ML Utils model architecture, optimiser, training settings, HGQ configuration, and optional hyperparameter search definitions.

For Optuna searches, direction must be explicitly defined as either maximize or minimize.

Example params.yml:

data: {}

meta:
  model_config: path/to/model.yml

  data_paths:
    train_path: /path/to/train.npz
    val_path: /path/to/val.npz
    test_path: /path/to/test.npz

  env_name: tf_v3
  save_path: out/

  search: optuna
  trials: 10
  sampler: bayesian
  objective: mean_auc
  direction: maximize

  cuts:
    min_score: 0.93
    max_ebops: 20000

Example model.yml:

inputs:
  jets:
    shape: [16]

layers:
  backbone:
    class: DenseLayers
    inputs: jets
    params:
      hidden_layer_sizes:
        type: list
        length: [1, 4]
        values: [16, 32, 64, 128]

      activation: relu
      use_hgq: true

  output:
    class: Dense
    inputs: backbone
    params:
      units: 5
      activation: softmax

outputs:
  output:
    loss: categorical_crossentropy

optimiser:
  class: Adam
  params:
    learning_rate:
      type: float
      values:
        - 0.01
        - 0.005
        - 0.001
        - 0.0005
        - 0.0001

training:
  batch_size:
    type: int
    values: [64, 128, 256, 512]

  epochs: 30

hgq_config:
  quantizer_scopes:
    - place: all
      default_q_type: kbi
      overflow_mode: SAT_SYM
      heterogeneous_axis: []

MLUtilsTrainer interprets the nested hyperparameter search definitions in model.yml and provides the sampled parameters to OptunaSearch.

OptunaSearch remains generic and also supports traditional flat search-space dictionaries used by other trainers.

If model.yml contains no hyperparameter search definitions, it is treated as a fixed configuration and the model is trained once without running Optuna.

The dataset paths are provided through meta.data_paths.

For quantization-aware training, Pareto checkpoints are produced during training and evaluated by Selector.

If selection cuts are provided under meta.cuts, the best checkpoint satisfying the requested constraints is selected.

If no cuts are provided, the best-scoring checkpoint from each trial is considered and the checkpoint with the lowest EBOP cost among those candidates is selected as the final model.

The selected model is reconstructed and finalized by MLUtilsFinalizer.

The final artifacts are stored as:

• winner.yml
• winner.keras

QAT callbacks and EBOP scheduling

For HGQ-enabled models, MLUtilsCallbackFactory builds the callbacks required for quantization-aware training.

HGQ is enabled at the model layer level with use_hgq: true. EBOP accounting and the initial penalty coefficient are configured through hgq_config:

layers:
  backbone:
    class: DenseLayers
    inputs: jets
    params:
      use_hgq: true

hgq_config:
  layer:
    enable_ebops: true
    beta0: 0.0

The evolution of the EBOP penalty during training is configured separately through qat.beta_schedule:

qat:
  beta_schedule:
    - [0, 0.0, linear]
    - [3, 1.0e-7, log]
    - [18, 1.0e-4, constant]

Each schedule entry specifies the epoch, target beta value, and interpolation mode. MLUtilsCallbackFactory converts this configuration into an HGQ PieceWiseSchedule and attaches it to training through BetaScheduler.

The callback factory is created and passed to MLUtilsTrainer:

from callbacks.multi_class_auc import MultiClassAUC
from wp21_train.callbacks.training_callbacks import MLUtilsCallbackFactory
from wp21_train.training.trainers.ml_utils_trainer import MLUtilsTrainer


def make_objective(valid_ds, resolved_config):
    output_node = next(iter(resolved_config["outputs"]))

    n_classes = resolved_config[
        "layers"
    ][output_node]["params"]["units"]

    return MultiClassAUC(
        valid_ds=valid_ds,
        output_node=output_node,
        n_classes=n_classes,
    )


callbacks_factory = MLUtilsCallbackFactory(
    objective_factory=make_objective,
    pareto_root="out/qat_pareto",
)

trainer = MLUtilsTrainer(
    params_config="params.yml",
    dataset=dataset,
    callbacks_factory=callbacks_factory,
)

For an HGQ-enabled model, the factory attaches the following callbacks:

• BetaScheduler, which applies the configured qat.beta_schedule;
• FreeEBOPs, which makes the model EBOP cost available during training;
• the user-provided objective callback;
• ParetoFront, which tracks the trade-off between predictive performance and EBOP cost and saves Pareto-optimal checkpoints.

Pareto checkpoints are stored separately for each trial:

out/qat_pareto/
  trial_0/
  trial_1/
  ...

Checkpoint filenames contain the epoch, objective score, and EBOP value:

epoch=<epoch>_score=<score>_ebops=<ebops>.weights.h5

The objective side of the Pareto front follows the configured optimization direction, while EBOPs are always minimized.

The saved Pareto checkpoints are later inspected by Selector, which applies the optional constraints from meta.cuts and chooses the final deployment candidate.

The MultiClassAUC is an example of a custom callable class for scoring a user should provide.

🧬 Type Utility

from wp21_train.utils.utility import get_short_type

print(get_short_type(42))     # 'd'
print(get_short_type("abc"))  # 's'

🧬 Logging Utility

from wp21_train.utils.logging import log_message

log_message("error", f"Provided number of events ({nevents}) is not an integer.")

🧪 Testing

pytest tests/

📜 Requirements

• Python ≥ 3.7
• uproot (for reading ROOT files)
• xml.etree.ElementTree (standard lib, for HLS/AIE parsing)
• CERN ROOT (installed and configured) if you use .root I/O
• PyYAML ≥ 6.0.0, < 7.0.0 (for reading yaml files)
• awkward (data is represented as awkward arrays)
• pybind11 and openmp - used to implement slow bits in C++

⚠️ Note About ROOT

This package supports .root file serialization and reading via CERN ROOT. If you intend to use this feature, ensure that ROOT is installed and properly sourced in your environment. You can install ROOT via Conda:

conda install -c conda-forge root

Or follow the official installation guide:
https://root.cern/install/

🧠 Versioning

The current package version is defined in:

from wp21_train.utils.version import __version__

🔖 License

This project is licensed under the MIT License. See the LICENSE file for details.

👤 Author

Ioannis Xiotidis
Email: ioannis.xiotidis@cern.ch

Pawel Mucha
Email: pawel.mucha@cern.ch

Vila Andela Petrovic
Email: vila.andela.petrovic@cern.ch

David Reikher
Email: david.reikher@cern.ch