stouputils.data_science.models.model_interface module#

Base implementation for machine learning models with common functionality. Provides shared infrastructure for model training, evaluation, and MLflow integration.

Implements core workflow methods:

  • Full training/evaluation pipeline (routine_full)

  • Transfer learning weight management

  • MLflow experiment tracking integration

class ModelInterface(num_classes: int, kfold: int = 0, transfer_learning: str = 'imagenet', **override_params: Any)[source]#

Bases: AbstractModel

Base class for all models containing common/public methods.

final_model: Any#

Attribute storing the final trained model (Keras model or PyTorch model).

model_name: str#

Attribute storing the name of the model class, automatically set from the class name. Used for logging and display purposes.

kfold: int#

Attribute storing the number of folds to use for K-fold cross validation. If 0 or 1, no K-fold cross validation is used. If > 1, uses K-fold cross validation with that many folds.

transfer_learning: str#

Attribute storing the transfer learning source, defaults to “imagenet”, can be set to None or a dataset name present in the data folder.

is_trained: bool#

Flag indicating if the model has been trained. Must be True before making predictions or evaluating the model.

num_classes: int#

Attribute storing the number of classes in the dataset.

override_params: dict[str, Any]#

Attribute storing the override parameters dictionary for the model.

run_name: str#

Attribute storing the name of the current run, automatically set during training.

history: list[dict[str, list[float]]]#

Attribute storing the training history for each fold.

evaluation_results: list[dict[str, float]]#

Attribute storing the evaluation results for each fold.

additional_training_data: XyTuple#

Attribute storing additional training data as a XyTuple that is incorporated into the training set right before model fitting.

This data bypasses cross-validation splitting and is only used during the training phase which differs from directly augmenting the dataset via dataset.training_data += additional_training_data, which would include the additional data in the cross-validation splitting process.

batch_size: int#

Attribute storing the batch size for training.

epochs: int#

Attribute storing the number of epochs for training.

class_weight: dict[int, float] | None#

0.34, 1: 0.66}.

Type:

Attribute storing the class weights for training, example

Type:

{0

unfreeze_percentage: float#

Attribute storing the percentage of layers to fine-tune from the last layer of the base model (0-100).

fine_tune_last_layers: int#

Attribute storing the number of layers to fine-tune, calculated from percentage when total_layers is known.

beta_1: float#

Attribute storing the beta 1 for Adam optimizer.

beta_2: float#

Attribute storing the beta 2 for Adam optimizer.

early_stop_patience: int#

Attribute storing the patience for early stopping.

learning_rate: float#

Attribute storing the learning rate for training.

reduce_lr_patience: int#

Attribute storing the patience for ReduceLROnPlateau.

min_delta: float#

Attribute storing the minimum delta for ReduceLROnPlateau (default of the library is 0.0001).

min_lr: float#

Attribute storing the minimum learning rate for ReduceLROnPlateau.

factor: float#

Attribute storing the factor for ReduceLROnPlateau.

warmup_epochs: int#

Attribute storing the number of epochs for learning rate warmup (0 to disable).

initial_warmup_lr: float#

Attribute storing the initial learning rate for warmup.

lr_finder_min_lr: float#

Attribute storing the minimum learning rate for the LR Finder.

lr_finder_max_lr: float#

Attribute storing the maximum learning rate for the LR Finder.

lr_finder_epochs: int#

Attribute storing the number of epochs for the LR Finder.

lr_finder_update_per_epoch: bool#

Attribute storing if the LR Finder should increase LR every epoch (True) or batch (False).

lr_finder_update_interval: int#

Attribute storing the number of steps between each lr increase, bigger value means more stable loss.

unfreeze_finder_epochs: int#

Attribute storing the number of epochs for the Unfreeze Percentage Finder

unfreeze_finder_update_per_epoch: bool#

Attribute storing if the Unfreeze Finder should unfreeze every epoch (True) or batch (False).

unfreeze_finder_update_interval: int#

Attribute storing the number of steps between each unfreeze, bigger value means more stable loss.

total_layers: int#

Attribute storing the total number of layers in the model.

class_load() → None[source]#

Clear histories, and set model parameters.

train(dataset: Dataset, verbose: int = 0) → bool[source]#

Method to train the model.

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data.

  • verbose (int) – Level of verbosity, decrease by 1 for each depth

Returns:

True if the model was trained successfully.

Return type:

bool

Raises:

ValueError – If the model could not be trained.

predict(X_test: Iterable[ndarray[Any, dtype[Any]]] | Dataset) → Iterable[ndarray[Any, dtype[Any]]][source]#

Method to predict the classes of a batch of data.

If a Dataset is provided, the test data ungrouped array will be used: X_test.test_data.ungrouped_array()[0]

Otherwise, the input is expected to be an Iterable of NDArray[Any].

Parameters:

X_test (Iterable[NDArray[Any]] | Dataset) – Features to use for prediction.

Returns:

Predictions of the batch.

Return type:

Iterable[NDArray[Any]]

Raises:

ValueError – If the model is not trained.

evaluate(dataset: Dataset, verbose: int = 0) → None[source]#

Method to evaluate the model, it will log metrics and plots to mlflow along with the model.

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data.

  • verbose (int) – Level of verbosity, decrease by 1 for each depth

routine_full(dataset: Dataset, verbose: int = 0) → ModelInterface[source]#

Method to perform a full routine (load, train and predict, evaluate, and export the model).

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data.

  • verbose (int) – Level of verbosity, decrease by 1 for each depth

Returns:

The model trained and evaluated.

Return type:

ModelInterface

_routine(dataset: Dataset, exp_name: str = '', verbose: int = 0)[source]#

Sub-method used in routine_full to perform a full routine

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data.

  • exp_name (str) – Name of the experiment (if empty, it will be set automatically)

  • verbose (int) – Level of verbosity, decrease by 1 for each depth

Returns:

The model trained and evaluated.

Return type:

ModelInterface

_get_transfer_learning_weights(dataset: Dataset, verbose: int = 0) → str[source]#

Get the transfer learning weights for the model.

This method handles retrieving pre-trained weights for transfer learning. It can:

  1. Return ‘imagenet’ weights for standard transfer learning

  2. Return None for no transfer learning

  3. Load weights from a previous training run on a different dataset

  4. Train a new model on a different dataset if no previous weights exist

Returns:

Path to weights file, ‘imagenet’, or None

Return type:

str

_get_total_layers() → int[source]#

Get the total number of layers in the model architecture, e.g. 427 for DenseNet121.

Compatible with Keras/TensorFlow and PyTorch models.

Returns:

Total number of layers in the model architecture.

Return type:

int

_set_parameters(override: dict[str, Any] | None = None) → None[source]#

Set some useful and common models parameters.

Parameters:

override (dict[str, Any]) – Dictionary of parameters to override.

_set_class_weight(y_train: ndarray[Any, dtype[Any]]) → None[source]#

Calculate class weight for balanced training.

Parameters:

y_train (NDArray[Any]) – Training labels

Returns:

Dictionary mapping class indices to weights, e.g. {0: 0.34, 1: 0.66}

Return type:

dict[int, float]

_log_parameters() → None[source]#

Log the model parameters.

_get_fold_split(training_data: XyTuple, kfold: int = 5) → Generator[tuple[XyTuple, XyTuple], None, None][source]#

Get fold split indices for cross validation.

This method splits the training data into k folds for cross validation while preserving the relationship between original images and their augmented versions.

The split is done using stratified k-fold to maintain class distribution across folds. For each fold, both the training and validation sets contain complete groups of original and augmented images.

Parameters:

  • training_data (XyTuple) – Dataset containing training and test data to split into folds

  • kfold (int) – Number of folds to create

Returns:

List of (train_data, val_data) tuples for each fold

Return type:

list[tuple[XyTuple, XyTuple]]

_train_final_model(dataset: Dataset, verbose: int = 0) → None[source]#

Train the final model on all data and return it.

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data

  • verbose (int) – Level of verbosity

_train_each_fold(dataset: Dataset, verbose: int = 0) → None[source]#

Train the model on each fold and fill self.models with the trained models.

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data

  • verbose (int) – Level of verbosity

class_train(dataset: Dataset, verbose: int = 0) → bool[source]#

Train the model using k-fold cross validation and then full model training.

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data (test data will not be used in class_train)

  • verbose (int) – Level of verbosity

Returns:

True if training was successful

Return type:

bool

_log_metrics(from_index: int = 1) → None[source]#

Calculate (average and standard deviation) and log metrics for each evaluation result.

Parameters:

from_index (int) – Index of the first evaluation result to use

class_evaluate(dataset: Dataset, metrics_names: tuple[str, ...] = (), save_model: bool = False, verbose: int = 0) → bool[source]#

Evaluate the model using the given predictions and labels.

Parameters:

  • dataset (Dataset) – Dataset containing the training and testing data

  • metrics_names (list[str]) – List of metrics to plot (default to all metrics)

  • save_model (bool) – Whether to save the best model

  • verbose (int) – Level of verbosity

Returns:

True if evaluation was successful

Return type:

bool

_find_best_learning_rate(dataset: Dataset, verbose: int = 0) → float[source]#

Find the best learning rate for the model, optionally using a subprocess.

Parameters:

  • dataset (Dataset) – Dataset to use for training.

  • verbose (int) – Verbosity level (controls progress bar).

Returns:

The best learning rate found.

Return type:

float

_find_best_unfreeze_percentage(dataset: Dataset, verbose: int = 0) → float[source]#

Find the best unfreeze percentage for the model, optionally using a subprocess.

Parameters:

  • dataset (Dataset) – Dataset to use for training.

  • verbose (int) – Verbosity level (controls progress bar).

Returns:

The best unfreeze percentage found.

Return type:

float

_train_fold(dataset: Dataset, fold_number: int = 0, mlflow_prefix: str = 'history', verbose: int = 0) → Any[source]#

Train model on a single fold.

Parameters:

  • dataset (Dataset) – Dataset to train on

  • fold_number (int) – Fold number (0 for final model)

  • prefix (str) – Prefix for the history

  • verbose (int) – Verbosity level

On this page

This Page