train_hgq

Functions:

• load_pretrained_model –

  Load a serialized Keras model from disk.

• main –

  Command-line interface for training a quantized (HGQ) profile reconstruction model.

• parse_yaml_config –

  Parse a YAML run configuration file with HGQ-specific custom tags.

• set_weights –

  Copy weights from a base model into another model, layer-by-layer.

• train –

  Train a quantized (HGQ) Keras model using the provided configuration and datasets.

load_pretrained_model

load_pretrained_model(model_path: Path) -> Model

Load a serialized Keras model from disk.

Parameters:

• model_path

  (Path) –

  Path to a saved Keras model (e.g., a .keras directory/file).

Returns:

• model ( Model ) –

  Deserialized Keras model instance.

Source code in src/fpga_profile_reco/core/train_hgq.py

def load_pretrained_model(model_path: Path) -> keras.Model:
    """
    Load a serialized Keras model from disk.

    Parameters
    ----------
    model_path : pathlib.Path
        Path to a saved Keras model (e.g., a ``.keras`` directory/file).

    Returns
    -------
    model : keras.Model
        Deserialized Keras model instance.
    """
    model = keras.models.load_model(model_path)
    return model

main

main()

Command-line interface for training a quantized (HGQ) profile reconstruction model.

This function parses command-line arguments, loads a YAML run configuration (with HGQ-specific custom tags), builds the training/validation datasets, instantiates :class:fpga_profile_reco.core.models.QHardNN, and runs training while writing logs and Pareto checkpoints to the configured output directories.

Command Line Parameters

• config : pathlib.Path Path to the YAML run configuration file.

Source code in src/fpga_profile_reco/core/train_hgq.py

def main():
    """
    Command-line interface for training a quantized (HGQ) profile reconstruction model.

    This function parses command-line arguments, loads a YAML run configuration
    (with HGQ-specific custom tags), builds the training/validation datasets,
    instantiates :class:`fpga_profile_reco.core.models.QHardNN`, and runs training
    while writing logs and Pareto checkpoints to the configured output directories.

    Command Line Parameters
    -----------------------
    - `config` : pathlib.Path
        Path to the YAML run configuration file.
    """
    import argparse
    import datetime
    import time

    from fpga_profile_reco.utils.helpers import format_time

    parser = argparse.ArgumentParser(description="Train a quantized model for equilibrium profile reconstruction.")
    parser.add_argument('config', type=Path, help="Path to the YAML run configuration file.")

    args = parser.parse_args()

    # set memory growth for GPUs
    gpus = tf.config.experimental.list_physical_devices('GPU')
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

    # read run configuration
    config = parse_yaml_config(args.config)

    train_ds, val_ds = get_datasets(train_radial_res=config['dataset']['train_radial_res'],
                                    val_radial_res=config['dataset']['val_radial_res'],
                                    batch_size=config['dataset']['batch_size'],
                                    rfp_only=config['dataset']['rfp_only'])

    # instantiate model
    model = QHardNN(architecture=config['architecture'], quantization=config['quantization'])
    # trigger model build to print summary
    model.build(input_shape=(None, 5))

    model.summary()

    start = time.time()
    print("\n\n===============================\n\n")
    print("Starting training at " + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    history = train(model=model, config=config, train_ds=train_ds, val_ds=val_ds)

    print("\n\n===============================\n\n")
    print("Trained terminated after ", len(history['loss']), "epochs.")
    print("Training finished at " + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    print(f"Total training time: {format_time(time.time() - start)}")

parse_yaml_config

parse_yaml_config(yaml_config_path: Path) -> dict

Parse a YAML run configuration file with HGQ-specific custom tags.

In addition to standard YAML types, this parser registers constructors on :class:yaml.SafeLoader for custom tags used in the quantized (HGQ) training pipeline, including constraints, quantizer configs, and scheduler objects.

Registered custom tags

• !tuple: build Python tuples
• !Min: instantiate :class:hgq.constraints.Min
• !Max: instantiate :class:hgq.constraints.Max
• !MinMax: instantiate :class:hgq.constraints.MinMax
• !QuantizerConfig: instantiate :class:hgq.config.QuantizerConfig
• !PieceWiseSchedule: instantiate :class:hgq.utils.sugar.PieceWiseSchedule
• !BetaScheduler: instantiate :class:hgq.utils.sugar.BetaScheduler
• !CosineAnnealingScheduler: instantiate :class:fpga_profile_reco.utils.schedulers.CosineAnnealingScheduler
• !CosineAnnealingWithRestartsScheduler: instantiate :class:fpga_profile_reco.utils.schedulers.CosineAnnealingWithRestartsScheduler

Notes

This function registers YAML constructors globally via :func:yaml.add_constructor (for yaml.SafeLoader).

Parameters:

• yaml_config_path

  (Path) –

  Path to the YAML configuration file.

Returns:

• config ( dict ) –

  Parsed configuration dictionary.

Source code in src/fpga_profile_reco/core/train_hgq.py

def parse_yaml_config(yaml_config_path: Path) -> dict:
    """
    Parse a YAML run configuration file with HGQ-specific custom tags.

    In addition to standard YAML types, this parser registers constructors on
    :class:`yaml.SafeLoader` for custom tags used in the quantized (HGQ) training
    pipeline, including constraints, quantizer configs, and scheduler objects.

    Registered custom tags
    ----------------------
    - ``!tuple``: build Python tuples
    - ``!Min``: instantiate :class:`hgq.constraints.Min`
    - ``!Max``: instantiate :class:`hgq.constraints.Max`
    - ``!MinMax``: instantiate :class:`hgq.constraints.MinMax`
    - ``!QuantizerConfig``: instantiate :class:`hgq.config.QuantizerConfig`
    - ``!PieceWiseSchedule``: instantiate :class:`hgq.utils.sugar.PieceWiseSchedule`
    - ``!BetaScheduler``: instantiate :class:`hgq.utils.sugar.BetaScheduler`
    - ``!CosineAnnealingScheduler``: instantiate
      :class:`fpga_profile_reco.utils.schedulers.CosineAnnealingScheduler`
    - ``!CosineAnnealingWithRestartsScheduler``: instantiate
      :class:`fpga_profile_reco.utils.schedulers.CosineAnnealingWithRestartsScheduler`

    Notes
    -----
    This function registers YAML constructors globally via
    :func:`yaml.add_constructor` (for ``yaml.SafeLoader``).

    Parameters
    ----------
    yaml_config_path : pathlib.Path
        Path to the YAML configuration file.

    Returns
    -------
    config : dict
        Parsed configuration dictionary.
    """
    # define constructor for custom object tags
    def construct_tuple(loader, node):
        return tuple(loader.construct_sequence(node))

    def construct_min(loader, node):
        mapping = loader.construct_mapping(node)
        return hgq.constraints.Min(min_value=mapping['min_value'])

    def construct_max(loader, node):
        mapping = loader.construct_mapping(node)
        return hgq.constraints.Max(max_value=mapping['max_value'])

    def construct_min_max(loader, node):
        mapping = loader.construct_mapping(node)
        return hgq.constraints.MinMax(min_value=mapping['min_value'], max_value=mapping['max_value'])

    def construct_quantizer_config(loader, node):
        mapping = loader.construct_mapping(node)
        return hgq.config.QuantizerConfig(**mapping)

    def construct_piecewise_schedule(loader, node):
        mapping = loader.construct_mapping(node, deep=True)
        return PieceWiseSchedule(intervals=mapping['intervals'])

    def construct_beta_scheduler(loader, node):
        mapping = loader.construct_mapping(node, deep=True)
        return BetaScheduler(beta_fn=mapping['beta_fn'])

    def construct_cosine_annealing_scheduler(loader, node):
        mapping = loader.construct_mapping(node)
        return CosineAnnealingScheduler(max_T=mapping['max_T'], min_lr=mapping['min_lr'])

    def construct_cosine_annealing_with_restarts_scheduler(loader, node):
        mapping = loader.construct_mapping(node, deep=True)
        return CosineAnnealingWithRestartsScheduler(restart_lrs=mapping['restart_lrs'], min_lrs=mapping['min_lrs'], Ts=mapping['Ts'])

    # register constructors
    yaml.add_constructor('!tuple', construct_tuple, Loader=yaml.SafeLoader)
    yaml.add_constructor('!Min', construct_min, Loader=yaml.SafeLoader)
    yaml.add_constructor('!Max', construct_max, Loader=yaml.SafeLoader)
    yaml.add_constructor('!MinMax', construct_min_max, Loader=yaml.SafeLoader)
    yaml.add_constructor('!QuantizerConfig', construct_quantizer_config, Loader=yaml.SafeLoader)
    yaml.add_constructor('!PieceWiseSchedule', construct_piecewise_schedule, Loader=yaml.SafeLoader)
    yaml.add_constructor('!BetaScheduler', construct_beta_scheduler, Loader=yaml.SafeLoader)
    yaml.add_constructor('!CosineAnnealingScheduler', construct_cosine_annealing_scheduler, Loader=yaml.SafeLoader)
    yaml.add_constructor('!CosineAnnealingWithRestartsScheduler', construct_cosine_annealing_with_restarts_scheduler, Loader=yaml.SafeLoader)

    # read yaml config file
    with open(yaml_config_path, 'r') as f:
        config = yaml.safe_load(f)

    return config

set_weights

set_weights(base_model: Model, model: Model) -> None

Copy weights from a base model into another model, layer-by-layer.

This function iterates over layers in base_model and model in lockstep (via :func:zip) and replaces the first two weight arrays of each target layer with those from the corresponding base layer.

Notes

• This assumes that corresponding layers have compatible weight structures and that the target layer has at least two weight tensors (commonly kernel and bias).
• Layers are matched purely by position, not by name.

Parameters:

• base_model

  (Model) –

  Model to copy weights from.

• model

  (Model) –

  Model to copy weights into.

Returns:

• None –

Source code in src/fpga_profile_reco/core/train_hgq.py

def set_weights(base_model: keras.Model, model: keras.Model) -> None:
    """
    Copy weights from a base model into another model, layer-by-layer.

    This function iterates over layers in ``base_model`` and ``model`` in lockstep
    (via :func:`zip`) and replaces the first two weight arrays of each target
    layer with those from the corresponding base layer.

    Notes
    -----
    - This assumes that corresponding layers have compatible weight structures
      and that the target layer has at least two weight tensors (commonly kernel
      and bias).
    - Layers are matched purely by position, not by name.

    Parameters
    ----------
    base_model : keras.Model
        Model to copy weights from.
    model : keras.Model
        Model to copy weights into.

    Returns
    -------
    None
    """
    for base_layer, layer in zip(base_model.layers, model.layers):
        weight_list = layer.get_weights()
        base_model_weights = base_layer.get_weights()
        # set weights and biases
        weight_list[0] = base_model_weights[0]
        weight_list[1] = base_model_weights[1]
        layer.set_weights(weight_list)

train

train(model: Model, config: dict, train_ds: Dataset, val_ds: Dataset) -> dict

Train a quantized (HGQ) Keras model using the provided configuration and datasets.

The model is compiled with an Adam optimizer using training.initial_lr. The loss/metrics are assumed to be handled by the model's internal/custom training logic. Training behavior is controlled by callbacks specified in the config, plus HGQ utilities such as EBOP accounting and Pareto checkpointing.

Parameters:

• model

  (Model) –

  Model to train.

• config

  (dict) –

  Run configuration dictionary as returned by :func:parse_yaml_config. Expected keys include run_config and training.

• train_ds

  (Dataset) –

  Training dataset.

• val_ds

  (Dataset) –

  Validation dataset.

Returns:

• history ( dict ) –

  History dictionary (i.e., history.history) returned by :meth:keras.Model.fit, mapping metric names to lists of epoch values.

Source code in src/fpga_profile_reco/core/train_hgq.py

def train(model: keras.Model, config: dict, train_ds: tf.data.Dataset, val_ds: tf.data.Dataset) -> dict:
    """
    Train a quantized (HGQ) Keras model using the provided configuration and datasets.

    The model is compiled with an Adam optimizer using ``training.initial_lr``.
    The loss/metrics are assumed to be handled by the model's internal/custom
    training logic. Training behavior is controlled by callbacks specified in
    the config, plus HGQ utilities such as EBOP accounting and Pareto checkpointing.

    Parameters
    ----------
    model : keras.Model
        Model to train.
    config : dict
        Run configuration dictionary as returned by :func:`parse_yaml_config`.
        Expected keys include ``run_config`` and ``training``.
    train_ds : tf.data.Dataset
        Training dataset.
    val_ds : tf.data.Dataset
        Validation dataset.

    Returns
    -------
    history : dict
        History dictionary (i.e., ``history.history``) returned by
        :meth:`keras.Model.fit`, mapping metric names to lists of epoch values.
    """
    run_config = config['run_config']
    training_config = config['training']

    # only compile with optimizer, loss and metrics are handled in the custom training loop
    model.compile(optimizer=keras.optimizers.Adam(learning_rate=training_config['initial_lr']))

    # load pretrained model weights if specified
    if run_config['load_pretrained_model']:
        print("Loading pretrained model weights from:", run_config['pretrained_model_path'])
        base_model = load_pretrained_model(run_config['pretrained_model_path'])
        set_weights(base_model, model)

    # setup various callbacks
    callbacks = []

    callbacks.append(keras.callbacks.TerminateOnNaN())
    if training_config['lr_scheduler']:
        callbacks.append(keras.callbacks.LearningRateScheduler(training_config['lr_scheduler'], verbose=1))
    if training_config['beta_scheduler']:
        callbacks.append(training_config['beta_scheduler'])
    callbacks.append(FreeEBOPs())
    tb_path = cfg.TENSORBOARD_LOGS_DIR / run_config['name']
    tb_path.mkdir(parents=True, exist_ok=True)
    callbacks.append(keras.callbacks.TensorBoard(log_dir=tb_path, histogram_freq=10, update_freq='epoch'))
    chkpt_path = cfg.PARETO_CHKPTS_DIR / run_config['name']
    chkpt_path.mkdir(parents=True, exist_ok=True)
    callbacks.append(ParetoFront(path=chkpt_path,
                                 fname_format='{epoch:04d}-val_loss-{val_loss:.4g}-ebops-{ebops:.4g}.keras',
                                 metrics=['val_loss', 'ebops'],
                                 enable_if=lambda x: x['val_loss'] < 1e-5,  # require a minimum val_loss to save
                                 sides=[-1, -1]))
    csv_path = cfg.HISTORY_DIR
    csv_path.mkdir(parents=True, exist_ok=True)
    callbacks.append(keras.callbacks.CSVLogger(filename=csv_path / (run_config['name'] + '.csv'), append=False))

    # run training
    history = model.fit(train_ds,
                        validation_data=val_ds,
                        callbacks=callbacks,
                        verbose=1,
                        epochs=training_config['epochs'])

    return history.history