Models Module

The models module provides base classes for machine learning models.

class xflow.models.BaseModel[source]

Bases: InferenceModel, Trainable, ABC

Combined abstract interface; implement a single subclass.

abstractmethod configure_optimizers() Any

Return optimizer(s)/schedulers required by the training framework.

abstractmethod classmethod load(path: str | PathLike[str], **kwargs) InferenceModel

Load model and config from disk.

abstractmethod predict(inputs: Any, **kwargs) Any

Run a forward/inference pass.

abstractmethod save(path: str | PathLike[str]) None

Persist weights (and any config) to disk.

set_train_mode(training: bool = True) None

Set model to training or evaluation mode. Override if needed.

abstractmethod training_step(batch: Tuple[Any, Any]) float | Dict[str, float | int | np.floating | np.integer]

Consume one batch (inputs, targets), perform an update, and return a loss or a metrics dict (if dict, must contain ‘loss’).

abstractmethod validation_step(batch: Tuple[Any, Any]) float | Dict[str, float | int | np.floating | np.integer]

Evaluate one batch in eval mode; return loss or metrics dict.