Training API

The training subsystem connects a validated experiment configuration to a model, data loaders, optimization components, metrics, loggers, and callbacks. Most applications configure these components in YAML; direct construction is useful for library integrations and tests. See the Training guide and configuration guide for workflows.

Configuration

`TrainConfig` contains runtime loop settings such as epochs, device, mixed precision, multi-loader sampling, scheduler stepping, and prediction caching.

chimera_ml.training.config.TrainConfig dataclass

Runtime configuration consumed by Trainer.

Attributes:

Name	Type	Description
epochs	int	Number of one-based training epochs.
grad_clip_norm	float | None	Optional maximum gradient norm; None disables clipping.
mixed_precision	bool	Enable autocast and gradient scaling on CUDA.
log_every_steps	int	Intended interval for progress logging.
device	str	Requested device string, typically cuda or cpu.
train_loader_mode	str	Multi-loader strategy: single, round_robin, or weighted.
train_stop_on	str	Whether a multi-loader epoch ends at first (min) or last (max) exhausted loader.
train_loader_weights	dict[str, float] | None	Optional per-loader weights for weighted mode.
use_scheduler	bool	Whether Trainer steps the configured scheduler.
scheduler_step_per_epoch	bool	Step once per epoch instead of after each optimizer update.
scheduler_monitor	str | None	Optional metric for metric-aware schedulers.
collect_cache	bool	Cache split predictions and metadata for callbacks.

Trainer

`Trainer` owns the fit/evaluate lifecycle. It moves batches to the selected device, calls the model and loss, updates metrics, logs values, invokes callbacks, and can cache split outputs for artifact-producing callbacks.

chimera_ml.training.trainer.Trainer dataclass

High-level fit/evaluate loop coordinating all training components.

The Trainer owns device placement, forward/loss execution, optimizer and scheduler stepping, metric accumulation, split normalization, logger integration, callback lifecycle, and optional prediction caching.

Attributes:

Name	Type	Description
model	BaseModel	Module receiving a Batch.
loss_fn	BaseLoss	Callable converting model output and batch into a loss tensor.
optimizer	Optimizer	Optimizer updated during fit.
metrics	list[BaseMetric]	Stateful metrics reset and computed around each epoch.
config	TrainConfig	Runtime TrainConfig.
mlflow_logger	BaseLogger | None	Optional logger with MLflow-style metrics/artifacts.
logger	Logger | None	Optional standard logger for console/file messages.
callbacks	list[BaseCallback]	Lifecycle extensions invoked by fit and evaluate.
scheduler	object | None	Optional learning-rate scheduler.
stop_training	bool	Flag checked between epochs; callbacks may set it.
global_step	int	Number of optimizer updates in the current fit.
cached_outputs	dict[str, CachedSplitOutputs]	Latest cached split outputs when enabled.

Note

fit resets global_step to zero. evaluate performs no optimizer updates but still runs callback hooks so evaluation callbacks can consume cached predictions.

fit

fit(train_loaders, val_loaders=None)

Run training for configured epochs and optional validation splits.

Parameters:

Name	Type	Description	Default
train_loaders	DataLoader | Mapping[str, DataLoader] | list[DataLoader] | tuple[DataLoader, ...]	One loader, a named mapping, or a sequence of training loaders.	required
val_loaders	DataLoader | Mapping[str, DataLoader] | list[DataLoader] | tuple[DataLoader, ...] | None	Optional loader container evaluated after each train epoch.	None

Raises: ValueError: If no training loader is provided or a training batch has no targets. Side Effects: Updates model parameters, optimizer/scheduler state, metrics, logger runs, callback state, and cached split outputs.

evaluate

evaluate(loaders, *, with_features=False, feature_extractor=None)

Evaluate loaders without optimization and return split metrics.

Parameters:

Name	Type	Description	Default
loaders	DataLoader | Mapping[str, DataLoader] | list[DataLoader] | tuple[DataLoader, ...]	One loader, a named mapping, or a sequence of loaders.	required
with_features	bool | None	Request feature extraction and include features in cached split outputs.	False
feature_extractor	FeatureExtractor | None	Optional callable used when with_features is true; otherwise model output aux features are used.	None

Returns: A mapping whose keys are prefixed with normalized split names. Raises: ValueError: If requested features are unavailable. Note: Callback hooks run in evaluation mode so validation-only callbacks can consume cached outputs.

get_cached_split_outputs

get_cached_split_outputs(split)

Return the latest cached outputs for a normalized split name.

Parameters:

Name	Type	Description	Default
split	str	Split key used by the Trainer.	required

Returns: Cached predictions, targets, features, and metadata, or None when caching was disabled or the split has not run.

Build context

`BuildContext` is shared while registry components are constructed. Components may read values, write derived values, and implement `describe_context` to enrich later builders.

chimera_ml.training.builders.BuildContext dataclass

Shared state for constructing one configured experiment.

Attributes:

Name	Type	Description
config	Any | None	Optional raw or typed experiment configuration.
stage	str | None	Optional construction stage label.
values	dict[str, Any]	Nested key/value store shared by builders.

register

register(component)

Let a built component publish derived values into the context.

Parameters:

Name	Type	Description	Default
component	Any	Object that may implement describe_context(context).	required

Returns: The same component, allowing use in builder expressions.

register_many

register_many(components)

Register each component and return the original list.

Parameters:

Name	Type	Description	Default
components	list[Any]	Built objects to inspect for describe_context.	required

Returns: The input list after context enrichment.

Registry builders

`build_from_registry` implements the common `name` + `params` configuration contract. The typed helpers inject runtime dependencies such as the model, optimizer, or context and are the usual bridge from YAML to Python objects.

chimera_ml.training.builders.build_from_registry

build_from_registry(registry, cfg, *, default_name=None, allow_none=False, normalize_name=True, params_key='params', name_key='name', inject=None, inject_overrides=True, smart_inject=False)

Universal builder for Registry-based factories.

Expected config format

cfg = {"name": "", "params": {...}}

Parameters:

Name	Type	Description	Default
registry	Any	Registry-like object with get(name) that returns a callable factory.	required
cfg	dict[str, Any] | None	Config dict or None.	required
default_name	str | None	Name used if cfg is None or cfg[name_key] is missing. If None and cfg is None: - return None if allow_none=True - else raise	None
allow_none	bool	If True, returns None when cfg is None (and default_name is None) or cfg has empty name.	False
normalize_name	bool	If True, uses lower() on name.	True
params_key	str	Config key containing factory parameters.	'params'
name_key	str	Config key containing the registry name.	'name'
inject	Mapping[str, Any] | None	Runtime dependencies (e.g. model_params, optimizer) merged into params.	None
inject_overrides	bool	If True, inject values override cfg params on conflicts.	True
smart_inject	bool	If True, only pass injected keys that are explicitly declared in the factory signature (based on signature introspection). Keys accepted only via **kwargs/**params are not injected implicitly. Useful when you want to inject runtime dependencies universally, but some factories should ignore them unless they opt in explicitly.	False

Returns:

Type	Description
Any	Built object (factory(**kwargs)) or None (if allow_none and cfg/default implies None).

chimera_ml.training.builders.build_datamodule

build_datamodule(cfg, *, context=None)

Build a data module from DATAMODULES configuration.

chimera_ml.training.builders.build_model

build_model(cfg, *, context=None)

Build a model from MODELS configuration.

chimera_ml.training.builders.build_loss

build_loss(cfg, *, context=None)

Build a loss from the LOSSES registry and optional context injection.

chimera_ml.training.builders.build_metrics

build_metrics(cfg_list, *, context=None)

Build metric instances from a list of registry configurations.

chimera_ml.training.builders.build_optimizer

build_optimizer(cfg, model, *, context=None)

Build an optimizer and inject the model into its factory.

When cfg is None, the registered adamw_optimizer factory is selected.

chimera_ml.training.builders.build_scheduler

build_scheduler(cfg, optimizer, *, context=None)

Build a scheduler, or return None when no scheduler is configured.

chimera_ml.training.builders.build_callbacks

build_callbacks(cfg_list, *, context=None)

Build callback instances from CALLBACKS configurations.

chimera_ml.training.builders.build_collate

build_collate(cfg, *, context=None)

Build a collate callable, defaulting to masking_collate.

chimera_ml.training.builders.build_logger

build_logger(cfg, *, inject=None, context=None)

Build optional logger with runtime injection.

Sweeps

`GridSweep` expands a finite Cartesian product or explicit trial list. `OptunaSweep` delegates typed suggestions and objective direction to Optuna. Both persist base/sweep/trial configurations and a manifest below the sweep log directory. `SweepTarget` defines the monitored metric and whether lower or higher values are preferred. See Sweeps.

chimera_ml.training.sweep.GridSweep

Bases: _SweepBase

Execute a finite grid sweep over explicit trials or parameter values.

The sweep accepts either a trials list or a parameters mapping. A parameters mapping expands to the Cartesian product, optionally truncated by max_trials. Each trial receives a copied base configuration and a persisted manifest record.

chimera_ml.training.sweep.OptunaSweep

Bases: _SweepBase

Execute an Optuna-backed sweep with typed parameter suggestions.

The parameters mapping describes Optuna search spaces, while SweepTarget defines the monitored log key and optimization direction. Study metadata, trial configs, and the best trial are persisted in the sweep manifest.

chimera_ml.utils.sweep.SweepTarget dataclass

Monitored log key and direction for sweep optimization.

Attributes:

Name	Type	Description
monitor	str	Scalar log key, commonly val/loss.
mode	str	Normalized direction, either min or max. minimize and maximize aliases are accepted in configuration.

from_config classmethod

from_config(sweep_cfg)

Create a target from target, objective, or metric config.