hyrax.verbs.train_stream

Attributes

logger

Classes

TrainStream	Streaming training verb — loads the model once, trains on batches on demand.
TrainStreamSession	Context manager for streaming training.

Module Contents

logger

class TrainStream(config)

Bases: hyrax.verbs.verb_registry.Verb

Streaming training verb — loads the model once, trains on batches on demand.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'train_stream'

add_parser_kwargs

description = 'Run streaming training: load the model once and train on batches interactively.'

REQUIRED_DATA_GROUPS = ('train_stream',)

OPTIONAL_DATA_GROUPS = ()

static setup_parser(parser)

No CLI arguments needed.

abstractmethod run_cli(args=None)

CLI stub — train_stream is a programmatic API only.

run() → TrainStreamSession

Set up the model and return a session for streaming training.

with hy.train_stream() as session:

for batch, metrics in session:

print(metrics[“loss”])

Because a live stream is open-ended (no epochs, no length, no train/validate/test splits), this verb does not use the Ignite create_trainer / trainer.run(..., max_epochs=...) path. Instead it reuses the same per-batch function the trainer’s engine calls (create_process_func() with "train_batch"), which already runs the model’s self-contained optimizer step, and drives it from the session loop.

Returns:

A context manager / session object. Iterate it (data-source driven) or call session.train_batch(batch) (manual); call session.close() when done to persist the final weights and get the trained model back.

Return type:

TrainStreamSession

Raises:

ValueError – If no [data_request.train_stream] is configured.

static _log_params(config, results_dir)

Log static run parameters to MLflow (mirrors Train._log_params).

class TrainStreamSession(process_func, model, config, results_dir, data_loader=None, provider=None)

Context manager for streaming training.

Holds a loaded model and its optimizer. When constructed with a data_loader (the data-source-driven path), the session is iterable and yields (batch, metrics) pairs as data arrives; otherwise feed batches yourself with train_batch().

Each processed batch runs the model’s self-contained training step (zero_grad → forward → loss → backward → optimizer.step) and returns its metrics dict. Model weights are persisted periodically (save_weights_every) and always on close().

Warning

train_batch() / process() are not thread-safe. Do not call them concurrently.

_process_func

_model

_config

_results_dir

data_loader = None

_provider = None

_closed = False

_batch_count = 0

_tb_logger

_close_logger

_best_loss = None

_context

__iter__()

Iterate the configured data source, training on each batch as it arrives.

Yields:

tuple[dict, dict | None] – The collated input batch and the training metrics for it (or None if the batch was skipped for being empty / smaller than min_batch_size).

Raises:

RuntimeError – If the session was created without a data source (no [data_request.train_stream] configuration).

_batch_num_samples(batch) → int | None

Return the number of samples in a collated batch, or None if unknown. batch should always be a dictionary at this point in the pipeline. And the key "object_id" is added by Hyrax to track the primary IDs of the samples, so it should always be present as well.

process(batch: dict) → dict | None

Run one training step on a single batch.

The batch is passed through the model’s train_batch (which performs the full optimizer step) unless it is empty or smaller than the configured min_batch_size, in which case it is skipped.

Parameters:

batch (dict) – Must contain "object_id" (list of str) and model-specific data fields.

Returns:

The training metrics (e.g. {"loss": ...}), or None if the batch was skipped.

Return type:

dict | None

Raises:

RuntimeError – If the session has already been closed.

train_batch

_log_metrics(result) → None

Log per-batch metrics to TensorBoard and (if a run is active) MLflow.

save_weights(file_name_suffix='') → None

Persist the current model weights to the results directory.

stop()

Signal the underlying streaming data source to stop iterating.

close()

Persist final weights, end logging, and return the trained model.

Returns:

The trained model.

Return type:

torch.nn.Module

checkpoint(model_metrics=None) → None

Save the current model weights if loss is lower than the previous best.

Parameters:

model_metrics (dict | None) – Dictionary containing model metrics. If None, the checkpoint is always saved.

__enter__()

__exit__(*_)