Trainer

Training orchestrator with checkpointing, evaluation, and optimizer lifecycle management.

class taktiny.trainer.Trainer(model, training_config, dataset_config, *, loss_fn, loss_has_aux=False, callbacks=None, compute_metrics=None)

Bases: TrainerEvaluateMixin, TrainerCheckpointMixin

Train a model on caller-provided batch iterables with an explicit loss.

DatasetConfig owns only batch placement/prefetch settings: dataset loading, preprocessing, batching, and sampling are external responsibilities. Training checkpoints use native Orbax, including dynamic model leaves, optional optimizer/EMA state, RNGs, and iterator progress. Static model configuration must be supplied again when constructing a resumed Trainer.

Parameters:

• model (Any)

• training_config (TrainingConfig)

• dataset_config (DatasetConfig)

• loss_fn (LossFn)

• loss_has_aux (bool)

• callbacks (Iterable[Any] | Any | None)

• compute_metrics (Callable[..., Mapping[str, Any]] | None)

add_callback(callback)

Append a callback and return it.

Return type:

Any

Parameters:

callback (Any)

remove_callback(callback)

Remove a previously registered callback.

Return type:

None

Parameters:

callback (Any)

extract_params()

Extract params based on the diagnosed model type.

Return type:

Any

train(resume_from_checkpoint=None)

Train on caller-provided batches; resume a native Orbax checkpoint.

Pass a checkpoint directory or ‘latest’. Supply the same model structure, optimizer, and data pipeline for exact resume. Async writes are drained on both success and failure before returning to the caller.

Return type:

None

Parameters:

resume_from_checkpoint (PathLike | None)

evaluate()

Evaluate the current model using validation_dataloader.

Return type:

dict[str, float]