stouputils 1.14.0__py3-none-any.whl

stouputils/data_science/models/model_interface.py

@@ -0,0 +1,939 @@
+""" Base implementation for machine learning models with common functionality.
+Provides shared infrastructure for model training, evaluation, and MLflow integration.
+
+Implements comprehensive workflow methods and features:
+
+Core Training & Evaluation:
+- Full training/evaluation pipeline (routine_full)
+- K-fold cross-validation with stratified splitting
+- Transfer learning weight management (ImageNet, custom datasets)
+- Model prediction and evaluation with comprehensive metrics
+
+Hyperparameter Optimization:
+- Learning Rate Finder with automatic best LR detection
+- Unfreeze Percentage Finder for fine-tuning optimization
+- Class weight balancing for imbalanced datasets
+- Learning rate warmup and scheduling (ReduceLROnPlateau)
+
+Advanced Training Features:
+- Early stopping with configurable patience
+- Model checkpointing with delay options
+- Additional training data integration (bypasses CV splitting)
+- Multi-processing support for memory management
+- Automatic retry mechanisms with error handling
+
+MLflow Integration:
+- Complete experiment tracking and logging
+- Parameter logging (training, optimizer, callback parameters)
+- Metric logging with averages and standard deviations
+- Model artifact saving and versioning
+- Training history visualization and plotting
+
+Model Architecture Support:
+- Keras/TensorFlow and PyTorch compatibility
+- Automatic layer counting and fine-tuning
+- Configurable unfreeze percentages for transfer learning
+- Memory leak prevention with subprocess training
+
+Evaluation & Visualization:
+- ROC and PR curve generation
+- Comprehensive metric calculation (Sensitivity, Specificity, AUC, etc.)
+- Training history plotting and analysis
+- Saliency maps and GradCAM visualization (single sample)
+- Cross-validation results aggregation
+
+Configuration & Utilities:
+- Extensive parameter override system
+- Verbosity control throughout pipeline
+- Temporary directory management for artifacts
+- Garbage collection and memory optimization
+- Error logging and handling with retry mechanisms
+"""
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownArgumentType=false
+
+# Imports
+from __future__ import annotations
+
+import gc
+import multiprocessing
+import multiprocessing.queues
+import time
+from collections.abc import Generator, Iterable
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import mlflow
+import numpy as np
+from mlflow.entities import Run
+from numpy.typing import NDArray
+from sklearn.utils import class_weight
+
+from ...decorators import handle_error, measure_time
+from ...print import progress, debug, info, warning
+from ...ctx import Muffle, MeasureTime
+from ...io import clean_path
+
+from .. import mlflow_utils
+from ..config.get import DataScienceConfig
+from ..dataset import Dataset, DatasetLoader, XyTuple
+from ..metric_dictionnary import MetricDictionnary
+from ..metric_utils import MetricUtils
+from ..utils import Utils
+from .abstract_model import AbstractModel
+
+# Constants
+MODEL_DOCSTRING: str = """ {model} implementation using advanced model class with common functionality.
+For information, refer to the ModelInterface class.
+"""
+CLASS_ROUTINE_DOCSTRING: str = """ Run the full routine for {model} model.
+
+Args:
+	dataset            (Dataset): Dataset to use for training and evaluation.
+	kfold              (int): K-fold cross validation index.
+	transfer_learning  (str): Pre-trained weights to use, can be "imagenet" or a dataset path like 'data/pizza_not_pizza'.
+	verbose            (int): Verbosity level.
+	**kwargs           (Any): Additional arguments.
+
+Returns:
+	{model}: Trained model instance.
+"""
+
+# Base class
+class ModelInterface(AbstractModel):
+	""" Base class for all models containing common/public methods. """
+
+	# Class constructor
+	def __init__(
+		self, num_classes: int, kfold: int = 0, transfer_learning: str = "imagenet", **override_params: Any
+	) -> None:
+		np.random.seed(DataScienceConfig.SEED)
+		multiprocessing.set_start_method("spawn", force=True)
+
+		## Base attributes
+		self.final_model: Any
+		""" Attribute storing the final trained model (Keras model or PyTorch model). """
+		self.model_name: str = self.__class__.__name__
+		""" Attribute storing the name of the model class, automatically set from the class name.
+		Used for logging and display purposes. """
+		self.kfold: int = kfold
+		""" 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. """
+		self.transfer_learning: str = transfer_learning
+		""" Attribute storing the transfer learning source, defaults to "imagenet",
+		can be set to None or a dataset name present in the data folder. """
+		self.is_trained: bool = False
+		""" Flag indicating if the model has been trained.
+		Must be True before making predictions or evaluating the model. """
+		self.num_classes: int = num_classes
+		""" Attribute storing the number of classes in the dataset. """
+		self.override_params: dict[str, Any] = override_params
+		""" Attribute storing the override parameters dictionary for the model. """
+		self.run_name: str = ""
+		""" Attribute storing the name of the current run, automatically set during training. """
+		self.history: list[dict[str, list[float]]] = []
+		""" Attribute storing the training history for each fold. """
+		self.evaluation_results: list[dict[str, float]] = []
+		""" Attribute storing the evaluation results for each fold. """
+		self.additional_training_data: XyTuple = XyTuple.empty()
+		""" 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.
+		"""
+
+
+		## Model parameters
+		# Training parameters
+		self.batch_size: int = 8
+		""" Attribute storing the batch size for training. """
+		self.epochs: int = 50
+		""" Attribute storing the number of epochs for training. """
+		self.class_weight: dict[int, float] | None = None
+		""" Attribute storing the class weights for training, e.g. {0: 0.34, 1: 0.66}. """
+
+		# Fine-tuning parameters
+		self.unfreeze_percentage: float = 100
+		""" Attribute storing the percentage of layers to fine-tune from the last layer of the base model (0-100). """
+		self.fine_tune_last_layers: int = -1
+		""" Attribute storing the number of layers to fine-tune, calculated from percentage when total_layers is known. """
+
+		# Optimizer parameters
+		self.beta_1: float = 0.95
+		""" Attribute storing the beta 1 for Adam optimizer. """
+		self.beta_2: float = 0.999
+		""" Attribute storing the beta 2 for Adam optimizer. """
+
+		# Callback parameters
+		self.early_stop_patience: int = 15
+		""" Attribute storing the patience for early stopping. """
+		self.model_checkpoint_delay: int = 0
+		""" Attribute storing the number of epochs before starting the checkpointing. """
+
+		# ReduceLROnPlateau parameters
+		self.learning_rate: float = 1e-4
+		""" Attribute storing the learning rate for training. """
+		self.reduce_lr_patience: int = 5
+		""" Attribute storing the patience for ReduceLROnPlateau. """
+		self.min_delta: float = 0.05
+		""" Attribute storing the minimum delta for ReduceLROnPlateau (default of the library is 0.0001). """
+		self.min_lr: float = 1e-7
+		""" Attribute storing the minimum learning rate for ReduceLROnPlateau. """
+		self.factor: float = 0.5
+		""" Attribute storing the factor for ReduceLROnPlateau. """
+
+		# Warmup parameters
+		self.warmup_epochs: int = 5
+		""" Attribute storing the number of epochs for learning rate warmup (0 to disable). """
+		self.initial_warmup_lr: float = 1e-7
+		""" Attribute storing the initial learning rate for warmup. """
+
+		# Learning Rate Finder parameters
+		self.lr_finder_min_lr: float = 1e-9
+		""" Attribute storing the *minimum* learning rate for the LR Finder. """
+		self.lr_finder_max_lr: float = 1.0
+		""" Attribute storing the *maximum* learning rate for the LR Finder. """
+		self.lr_finder_epochs: int = 3
+		""" Attribute storing the number of epochs for the LR Finder. """
+		self.lr_finder_update_per_epoch: bool = False
+		""" Attribute storing if the LR Finder should increase LR every epoch (True) or batch (False). """
+		self.lr_finder_update_interval: int = 5
+		""" Attribute storing the number of steps between each lr increase, bigger value means more stable loss. """
+
+		# Unfreeze Percentage Finder parameters
+		self.unfreeze_finder_epochs: int = 500
+		""" Attribute storing the number of epochs for the Unfreeze Percentage Finder """
+		self.unfreeze_finder_update_per_epoch: bool = True
+		""" Attribute storing if the Unfreeze Finder should unfreeze every epoch (True) or batch (False). """
+		self.unfreeze_finder_update_interval: int = 25
+		""" Attribute storing the number of steps between each unfreeze, bigger value means more stable loss. """
+
+		## Model architecture
+		self.total_layers: int = 0
+		""" Attribute storing the total number of layers in the model. """
+
+	# String representation
+	def __str__(self) -> str:
+		return f"{self.model_name} (is_trained: {self.is_trained})"
+
+	# Public methods
+	@classmethod
+	def class_routine(
+		cls, dataset: Dataset, kfold: int = 0, transfer_learning: str = "imagenet", verbose: int = 0, **override_params: Any
+	) -> ModelInterface:
+		return cls(dataset.num_classes, kfold, transfer_learning, **override_params).routine_full(dataset, verbose)
+
+	@measure_time(printer=debug, message="Class load (ModelInterface)")
+	def class_load(self) -> None:
+		""" Clear histories, and set model parameters. """
+		# Initialize some attributes
+		self.history.clear()
+		self.evaluation_results.clear()
+
+		# Get the total number of layers in a subprocess to avoid memory leaks with tensorflow
+		with multiprocessing.Pool(1) as pool:
+			self.total_layers = pool.apply(self._get_total_layers)
+
+		# Create final model by connecting input to output layer
+		self._set_parameters(self.override_params)
+
+	@measure_time
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def train(self, dataset: Dataset, verbose: int = 0) -> bool:
+		""" Method to train the model.
+
+		Args:
+			dataset  (Dataset):  Dataset containing the training and testing data.
+			verbose  (int):      Level of verbosity, decrease by 1 for each depth
+		Returns:
+			bool: True if the model was trained successfully.
+		Raises:
+			ValueError: If the model could not be trained.
+		"""
+		if not self.class_train(dataset, verbose=verbose):
+			raise ValueError("The model could not be trained.")
+		self.is_trained = True
+		return True
+
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def predict(self, X_test: Iterable[NDArray[Any]] | Dataset) -> Iterable[NDArray[Any]]:
+		""" 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].
+
+		Args:
+			X_test  (Iterable[NDArray[Any]] | Dataset):     Features to use for prediction.
+		Returns:
+			Iterable[NDArray[Any]]: Predictions of the batch.
+		Raises:
+			ValueError: If the model is not trained.
+		"""
+		if not self.is_trained:
+			raise ValueError("The model must be trained before predicting.")
+
+		# Get X_test from Dataset
+		if isinstance(X_test, Dataset):
+			return self.class_predict(X_test.test_data.ungrouped_array()[0])
+		else:
+			return self.class_predict(X_test)
+
+
+	@measure_time
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def evaluate(self, dataset: Dataset, verbose: int = 0) -> None:
+		""" Method to evaluate the model, it will log metrics and plots to mlflow along with the model.
+
+		Args:
+			dataset  (Dataset):  Dataset containing the training and testing data.
+			verbose  (int):      Level of verbosity, decrease by 1 for each depth
+		"""
+		if not self.is_trained:
+			raise ValueError("The model must be trained before evaluating.")
+
+		# Metrics (Sensibility, Specificity, AUC, etc.)
+		predictions: Iterable[NDArray[Any]] = self.predict(dataset)
+		metrics: dict[str, float] = MetricUtils.metrics(dataset, predictions, self.run_name)
+		mlflow.log_metrics(metrics)
+
+		# Model specific evaluation
+		self.class_evaluate(dataset, save_model=DataScienceConfig.SAVE_MODEL, verbose=verbose)
+
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	@measure_time
+	def routine_full(self, dataset: Dataset, verbose: int = 0) -> ModelInterface:
+		""" Method to perform a full routine (load, train and predict, evaluate, and export the model).
+
+		Args:
+			dataset  (Dataset):  Dataset containing the training and testing data.
+			verbose  (int):      Level of verbosity, decrease by 1 for each depth
+		Returns:
+			ModelInterface: The model trained and evaluated.
+		"""
+		# Get the transfer learning weights
+		self.transfer_learning = self._get_transfer_learning_weights(dataset, verbose=verbose)
+
+		# Perform the routine
+		return self._routine(dataset, verbose=verbose)
+
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def _routine(self, dataset: Dataset, exp_name: str = "", verbose: int = 0):
+		""" Sub-method used in routine_full to perform a full routine
+
+		Args:
+			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:
+			ModelInterface: The model trained and evaluated.
+		"""
+		# Init the model
+		self.class_load()
+
+		# Start mlflow run silently
+		with MeasureTime(message="Experiment setup time"):
+			with Muffle(mute_stderr=True):
+				exp_name = dataset.get_experiment_name() if exp_name == "" else exp_name
+				self.run_name = mlflow_utils.start_run(DataScienceConfig.MLFLOW_URI, exp_name, self.model_name)
+
+		# Log the dataset used for data augmentation and parameters
+		if dataset.original_dataset:
+			mlflow.log_params({"data_augmentation_based_of": dataset.original_dataset.name})
+		self._log_parameters()
+
+		# Train the model
+		self.train(dataset, verbose)
+
+		# Evaluate the model
+		self.evaluate(dataset, verbose)
+
+		# End mlflow run and return the model
+		with Muffle(mute_stderr=True):
+			mlflow.end_run()
+		return self
+
+
+
+	# Protected methods
+	def _get_transfer_learning_weights(self, dataset: Dataset, verbose: int = 0) -> str:
+		""" 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:
+			str: Path to weights file, 'imagenet', or None
+		"""
+		# If transfer is None or imagenet, return it
+		if self.transfer_learning in ("imagenet", None):
+			return self.transfer_learning
+
+		# Else, find the weights file path
+		else:
+			dataset_name: str = clean_path(self.transfer_learning).split("/")[-1]
+			exp_name: str = dataset.get_experiment_name(override_name=dataset_name)
+
+			# Find a run with the same model name, and get the weights file
+			with Muffle(mute_stderr=True):
+				runs: list[Run] = mlflow_utils.get_runs_by_model_name(exp_name, self.model_name)
+
+			# If no runs are found, train a new model on the dataset
+			if len(runs) == 0:
+
+				# Load dataset
+				pre_dataset: Dataset = DatasetLoader.from_path(
+					self.transfer_learning,
+					loading_type=dataset.loading_type,
+					grouping_strategy=dataset.grouping_strategy
+				)
+				info(f"In order to do Transfer Learning, training the model on the dataset '{pre_dataset}' first.")
+
+				# Save current settings
+				previous_transfer_learning: str = self.transfer_learning
+				previous_save_model: bool = DataScienceConfig.SAVE_MODEL
+				previous_kfold: int = self.kfold
+
+				# Configure for transfer learning training
+				self.transfer_learning = "imagenet"  # Start with imagenet weights
+				DataScienceConfig.SAVE_MODEL = True  # Enable model saving
+				self.kfold = 0         # Disable k-fold
+
+				# Train model on transfer learning dataset
+				self._routine(pre_dataset, exp_name=exp_name, verbose=verbose)
+
+				# Restore previous settings
+				self.transfer_learning = previous_transfer_learning
+				DataScienceConfig.SAVE_MODEL = previous_save_model
+				self.kfold = previous_kfold
+
+				# Get the weights file path - need to refresh the experiment object
+				runs: list[Run] = mlflow_utils.get_runs_by_model_name(exp_name, self.model_name)
+
+			# If no runs are found, raise an error
+			if not runs:
+				raise ValueError(f"No runs found for model {self.model_name} in experiment {exp_name}")
+			run: Run = runs[-1]
+
+			# Get the last run's weights path
+			# FIXME: Only works if MLFLow URI is file-tree based (not remote or sqlite), which is default
+			return mlflow_utils.get_weights_path(from_string=str(run.info.artifact_uri))
+
+	def _get_total_layers(self) -> int:
+		""" Get the total number of layers in the model architecture, e.g. 427 for DenseNet121.
+
+		Compatible with Keras/TensorFlow and PyTorch models.
+
+		Returns:
+			int: Total number of layers in the model architecture.
+		"""
+		architecture: Any = self._get_architectures()[1]
+		total_layers: int = 0
+		# Keras/TensorFlow
+		if hasattr(architecture, "layers"):
+			total_layers = len(architecture.layers)
+
+		# PyTorch
+		elif hasattr(architecture, "children"):
+			total_layers = len(architecture.children())
+
+		# Free memory and return the total number of layers
+		del architecture
+		gc.collect()
+		return total_layers
+
+	def _set_parameters(self, override: dict[str, Any] | None = None) -> None:
+		""" Set some useful and common models parameters.
+
+		Args:
+			override (dict[str, Any]): Dictionary of parameters to override.
+		"""
+		if override is None:
+			override = {}
+
+		# Training parameters
+		self.batch_size = override.get("batch_size", self.batch_size)
+		self.epochs = override.get("epochs", self.epochs)
+
+		# Callback parameters
+		self.early_stop_patience = override.get("early_stop_patience", self.early_stop_patience)
+		self.model_checkpoint_delay = override.get("model_checkpoint_delay", self.model_checkpoint_delay)
+
+		# ReduceLROnPlateau parameters
+		self.learning_rate = override.get("learning_rate", self.learning_rate)
+		self.reduce_lr_patience = override.get("reduce_lr_patience", self.reduce_lr_patience)
+		self.min_delta = override.get("min_delta", self.min_delta)
+		self.min_lr = override.get("min_lr", self.min_lr)
+		self.factor = override.get("factor", self.factor)
+
+		# Warmup parameters
+		self.warmup_epochs = override.get("warmup_epochs", self.warmup_epochs)
+		self.initial_warmup_lr = override.get("initial_warmup_lr", self.initial_warmup_lr)
+
+		# Fine-tune parameters
+		self.unfreeze_percentage = override.get("unfreeze_percentage", self.unfreeze_percentage)
+		self.fine_tune_last_layers = max(1, round(self.total_layers * self.unfreeze_percentage / 100))
+
+		# Optimizer parameters
+		self.beta_1 = override.get("beta_1", self.beta_1)
+		self.beta_2 = override.get("beta_2", self.beta_2)
+
+		# Learning Rate Finder parameters
+		self.lr_finder_min_lr = override.get("lr_finder_min_lr", self.lr_finder_min_lr)
+		self.lr_finder_max_lr = override.get("lr_finder_max_lr", self.lr_finder_max_lr)
+		self.lr_finder_epochs = override.get("lr_finder_epochs", self.lr_finder_epochs)
+		self.lr_finder_update_per_epoch = override.get("lr_finder_update_per_epoch", self.lr_finder_update_per_epoch)
+		self.lr_finder_update_interval = override.get("lr_finder_update_interval", self.lr_finder_update_interval)
+
+		# Unfreeze Percentage Finder parameters
+		self.unfreeze_finder_epochs = override.get("unfreeze_finder_epochs", self.unfreeze_finder_epochs)
+		self.unfreeze_finder_update_per_epoch = override.get("unfreeze_finder_update_per_epoch", self.unfreeze_finder_update_per_epoch)
+		self.unfreeze_finder_update_interval = override.get("unfreeze_finder_update_interval", self.unfreeze_finder_update_interval)
+
+		# Other parameters
+		self.additional_training_data += override.get("additional_training_data", XyTuple.empty())
+
+
+	def _set_class_weight(self, y_train: NDArray[Any]) -> None:
+		""" Calculate class weight for balanced training.
+
+		Args:
+			y_train (NDArray[Any]): Training labels
+		Returns:
+			dict[int, float]: Dictionary mapping class indices to weights, e.g. {0: 0.34, 1: 0.66}
+		"""
+		# Get the true classes (one-hot -> class indices)
+		true_classes: NDArray[Any] = Utils.convert_to_class_indices(y_train)
+
+		# Set the class weights (balanced)
+		self.class_weight = dict(enumerate(class_weight.compute_class_weight(
+			class_weight="balanced",
+			classes=np.unique(true_classes),
+			y=true_classes
+		)))
+
+	def _log_parameters(self) -> None:
+		""" Log the model parameters. """
+		mlflow.log_params({
+			"cfg_test_size": DataScienceConfig.TEST_SIZE,
+			"cfg_validation_size": DataScienceConfig.VALIDATION_SIZE,
+			"cfg_seed": DataScienceConfig.SEED,
+			"cfg_save_model": DataScienceConfig.SAVE_MODEL,
+			"cfg_device": DataScienceConfig.TENSORFLOW_DEVICE,
+
+			# Base attributes
+			"param_kfold": self.kfold,
+			"param_transfer_learning": self.transfer_learning,
+
+			# Training parameters
+			"param_batch_size": self.batch_size,
+			"param_epochs": self.epochs,
+
+			# Fine-tuning parameters
+			"param_unfreeze_percentage": self.unfreeze_percentage,
+			"param_fine_tune_last_layers": self.fine_tune_last_layers,
+			"param_total_layers": self.total_layers,
+
+			# Optimizer parameters
+			"param_beta_1": self.beta_1,
+			"param_beta_2": self.beta_2,
+			"param_learning_rate": self.learning_rate,
+
+			# Callback parameters
+			"param_early_stop_patience": self.early_stop_patience,
+			"param_model_checkpoint_delay": self.model_checkpoint_delay,
+
+			# ReduceLROnPlateau parameters
+			"param_reduce_lr_patience": self.reduce_lr_patience,
+			"param_min_delta": self.min_delta,
+			"param_min_lr": self.min_lr,
+			"param_factor": self.factor,
+
+			# Warmup parameters
+			"param_warmup_epochs": self.warmup_epochs,
+			"param_initial_warmup_lr": self.initial_warmup_lr,
+		})
+
+	def _get_fold_split(self, training_data: XyTuple, kfold: int = 5) -> Generator[tuple[XyTuple, XyTuple], None, None]:
+		""" 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.
+
+		Args:
+			training_data  (XyTuple): Dataset containing training and test data to split into folds
+			kfold          (int):     Number of folds to create
+		Returns:
+			list[tuple[XyTuple, XyTuple]]: List of (train_data, val_data) tuples for each fold
+		"""
+		assert kfold not in (0, 1), "kfold must not be 0 or 1"
+		yield from training_data.kfold_split(n_splits=kfold, random_state=DataScienceConfig.SEED)
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def _train_final_model(self, dataset: Dataset, verbose: int = 0) -> None:
+		""" Train the final model on all data and return it.
+
+		Args:
+			dataset  (Dataset):  Dataset containing the training and testing data
+			verbose  (int):      Level of verbosity
+		"""
+		# Get validation data from training data
+		debug(f"Training final model on train/val split: {dataset}")
+
+		# Verbose info message
+		if verbose > 0:
+			info(
+				f"({self.model_name}) Training final model on full dataset with "
+				f"{len(dataset.training_data.X)} samples ({len(dataset.val_data.X)} validation)"
+			)
+
+		# Put the validation data in the test data (since we don't use the test data in the train function)
+		old_test_data: XyTuple = dataset.test_data
+		dataset.test_data = dataset.val_data
+
+		# Train the final model and remember it
+		self.final_model = self._train_fold(dataset, fold_number=0, mlflow_prefix="history_final", verbose=verbose)
+
+		# Restore the old test data
+		dataset.test_data = old_test_data
+		gc.collect()
+
+	@measure_time
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def _train_each_fold(self, dataset: Dataset, verbose: int = 0) -> None:
+		""" Train the model on each fold and fill self.models with the trained models.
+
+		Args:
+			dataset  (Dataset):  Dataset containing the training and testing data
+			verbose  (int):      Level of verbosity
+		"""
+		# Get fold split
+		fold_split: Generator[tuple[XyTuple, XyTuple], None, None] = self._get_fold_split(dataset.training_data, self.kfold)
+
+		# Train on each fold
+		for i, (train_data, test_data) in enumerate(fold_split):
+			fold_number: int = i + 1
+
+			# During Cross Validation, the validation data is the same as the test data.
+			# Except when the validation population is 1 sample (e.g. LeaveOneOut)
+			# Therefore, we need to use the original validation data for the final model
+			if self.kfold < 0 or len(test_data.X) == 1:
+				val_data: XyTuple = dataset.val_data
+			else:
+				val_data: XyTuple = test_data
+
+			# Create a new dataset (train/val based of training data)
+			new_dataset: Dataset = Dataset(
+				training_data=train_data,
+				val_data=val_data,
+				test_data=test_data,
+				name=dataset.name,
+				grouping_strategy=dataset.grouping_strategy,
+				labels=dataset.labels
+			)
+
+			# Log the fold
+			if verbose > 0:
+				# If there are multiple validation samples or no filepaths, show the number of validation samples
+				if len(test_data.X) != 1 or not test_data.filepaths:
+					debug(
+						f"({self.model_name}) Fold {fold_number} training with "
+						f"{len(train_data.X)} samples ({len(test_data.X)} validation)"
+					)
+				# Else, show the filepath of the single validation sample (useful for debugging)
+				else:
+					debug(
+						f"({self.model_name}) Fold {fold_number} training with "
+						f"{len(train_data.X)} samples (validation: {test_data.filepaths[0]})"
+					)
+
+			# Train the model on the fold
+			handle_error(self._train_fold,
+				message=f"({self.model_name}) Fold {fold_number} training failed", error_log=DataScienceConfig.ERROR_LOG
+			)(
+				dataset=new_dataset,
+				fold_number=fold_number,
+				mlflow_prefix=f"history_fold_{fold_number}",
+				verbose=verbose
+			)
+
+			# Collect garbage to free up some memory
+			gc.collect()
+
+	@measure_time
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def class_train(self, dataset: Dataset, verbose: int = 0) -> bool:
+		""" Train the model using k-fold cross validation and then full model training.
+
+		Args:
+			dataset (Dataset):  Dataset containing the training and testing data (test data will not be used in class_train)
+			verbose (int):      Level of verbosity
+		Returns:
+			bool: True if training was successful
+		"""
+		# Compute the class weights
+		self._set_class_weight(np.array(dataset.training_data.y))
+
+		# Find the best learning rate
+		if DataScienceConfig.DO_LEARNING_RATE_FINDER > 0:
+			info(f"({self.model_name}) Finding the best learning rate...")
+			found_lr: float | None = self._find_best_learning_rate(dataset, verbose)
+			if DataScienceConfig.DO_LEARNING_RATE_FINDER == 2 and found_lr is not None:
+				self.learning_rate = found_lr
+				mlflow.log_params({"param_learning_rate": found_lr})
+				info(f"({self.model_name}) Now using learning rate: {found_lr:.2e}")
+
+		# Find the best unfreeze percentage
+		if DataScienceConfig.DO_UNFREEZE_FINDER > 0:
+			info(f"({self.model_name}) Finding the best unfreeze percentage...")
+			found_unfreeze: float | None = self._find_best_unfreeze_percentage(dataset, verbose)
+			if DataScienceConfig.DO_UNFREEZE_FINDER == 2 and found_unfreeze is not None:
+				self.unfreeze_percentage = found_unfreeze
+				info(f"({self.model_name}) Now using unfreeze percentage: {found_unfreeze:.2f}%")
+
+		# If k-fold is enabled, train the model on each fold
+		if self.kfold not in (0, 1):
+			self._train_each_fold(dataset, verbose)
+
+		# Train the final model on all data
+		self._train_final_model(dataset, verbose)
+		return True
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def _log_metrics(self, from_index: int = 1) -> None:
+		""" Calculate (average and standard deviation) and log metrics for each evaluation result.
+
+		Args:
+			from_index (int): Index of the first evaluation result to use
+		"""
+		# For each metric, calculate the average and standard deviation
+		for metric_name in self.evaluation_results[0].keys():
+
+			# Get the metric values for each fold
+			metric_values: list[float] = [x[metric_name] for x in self.evaluation_results[from_index:]]
+			if not metric_values:
+				continue
+
+			# Log the average and standard deviation
+			avg_key: str = MetricDictionnary.AVERAGE_METRIC.replace("METRIC_NAME", metric_name)
+			std_key: str = MetricDictionnary.STANDARD_DEVIATION_METRIC.replace("METRIC_NAME", metric_name)
+			mlflow.log_metric(avg_key, float(np.mean(metric_values)))
+			mlflow.log_metric(std_key, float(np.std(metric_values)))
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def class_evaluate(
+		self,
+		dataset: Dataset,
+		metrics_names: tuple[str, ...] = (),
+		save_model: bool = False,
+		verbose: int = 0
+	) -> bool:
+		""" Evaluate the model using the given predictions and labels.
+
+		Args:
+			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:
+			bool: True if evaluation was successful
+		"""
+		# If no metrics names are provided, use all metrics
+		if not metrics_names:
+			metrics_names = tuple(self.evaluation_results[0].keys())
+
+		# Log metrics and plot curves
+		MetricUtils.plot_every_metric_curves(self.history, metrics_names, self.run_name)
+		self._log_metrics()
+
+		# Save the best model if save_model is True
+		if save_model:
+			if verbose > 0:
+				with MeasureTime(debug, "Saving best model"):
+					self._log_final_model()
+			else:
+				self._log_final_model()
+
+		# Success
+		return True
+
+
+	# Protected methods for training
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def _find_best_learning_rate(self, dataset: Dataset, verbose: int = 0) -> float:
+		""" Find the best learning rate for the model, optionally using a subprocess.
+
+		Args:
+			dataset  (Dataset):  Dataset to use for training.
+			verbose  (int):      Verbosity level (controls progress bar).
+		Returns:
+			float: The best learning rate found.
+		"""
+		results: dict[str, Any] = {}
+		for try_count in range(10):
+			try:
+				if DataScienceConfig.DO_FIT_IN_SUBPROCESS:
+					queue: multiprocessing.queues.Queue[dict[str, Any]] = multiprocessing.Queue()
+					process: multiprocessing.Process = multiprocessing.Process(
+						target=self._find_best_learning_rate_subprocess,
+						kwargs={"dataset": dataset, "queue": queue, "verbose": verbose}
+					)
+					process.start()
+					process.join()
+					results = queue.get(timeout=60)
+				else:
+					results = self._find_best_learning_rate_subprocess(dataset, verbose=verbose)
+				if results:
+					break
+			except Exception as e:
+				warning(f"Error finding best learning rate: {e}\nRetrying in 60 seconds ({try_count + 1}/10)...")
+				time.sleep(60)
+
+		# Plot the learning rate vs loss and find the best learning rate
+		return MetricUtils.find_best_x_and_plot(
+			results["learning_rates"],
+			results["losses"],
+			smoothen=True,
+			use_steep=True,
+			run_name=self.run_name,
+			x_label="Learning Rate",
+			y_label="Loss",
+			plot_title="Learning Rate Finder",
+			log_x=True,
+			y_limits=(0, 4.0)
+		)
+
+	@handle_error(error_log=DataScienceConfig.ERROR_LOG)
+	def _find_best_unfreeze_percentage(self, dataset: Dataset, verbose: int = 0) -> float:
+		""" Find the best unfreeze percentage for the model, optionally using a subprocess.
+
+		Args:
+			dataset  (Dataset):  Dataset to use for training.
+			verbose  (int):      Verbosity level (controls progress bar).
+		Returns:
+			float: The best unfreeze percentage found.
+		"""
+		results: dict[str, Any] = {}
+		for try_count in range(10):
+			try:
+				if DataScienceConfig.DO_FIT_IN_SUBPROCESS:
+					queue: multiprocessing.queues.Queue[dict[str, Any]] = multiprocessing.Queue()
+					process: multiprocessing.Process = multiprocessing.Process(
+						target=self._find_best_unfreeze_percentage_subprocess,
+						kwargs={"dataset": dataset, "queue": queue, "verbose": verbose}
+					)
+					process.start()
+					process.join()
+					results = queue.get(timeout=60)
+				else:
+					results = self._find_best_unfreeze_percentage_subprocess(dataset, verbose=verbose)
+				if results:
+					break
+			except Exception as e:
+				warning(f"Error finding best unfreeze percentage: {e}\nRetrying in 60 seconds ({try_count + 1}/10)...")
+				time.sleep(60)
+
+		# Plot the unfreeze percentage vs loss and find the best unfreeze percentage
+		return MetricUtils.find_best_x_and_plot(
+			results["unfreeze_percentages"],
+			results["losses"],
+			smoothen=True,
+			use_steep=False,
+			run_name=self.run_name,
+			x_label="Unfreeze Percentage",
+			y_label="Loss",
+			plot_title="Unfreeze Percentage Finder",
+			log_x=False,
+			y_limits=(0, 4.0)
+		)
+
+
+	@measure_time
+	def _train_fold(self, dataset: Dataset, fold_number: int = 0, mlflow_prefix: str = "history", verbose: int = 0) -> Any:
+		""" Train model on a single fold.
+
+		Args:
+			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
+		"""
+		# Create the checkpoint path
+		checkpoint_path: str = f"{DataScienceConfig.TEMP_FOLDER}/{self.run_name}_best_model_fold_{fold_number}.keras"
+
+		# Prepare visualization arguments if needed
+		temp_dir: TemporaryDirectory[str] | None = None
+		if DataScienceConfig.DO_SALIENCY_AND_GRADCAM and dataset.test_data.n_samples == 1:
+			temp_dir = TemporaryDirectory()
+
+		# Create and run the process
+		return_values: dict[str, Any] = {}
+		for try_count in range(10):
+			try:
+				if DataScienceConfig.DO_FIT_IN_SUBPROCESS and fold_number > 0:
+					queue: multiprocessing.queues.Queue[dict[str, Any]] = multiprocessing.Queue()
+					process: multiprocessing.Process = multiprocessing.Process(
+						target=self._train_subprocess,
+						args=(dataset, checkpoint_path, temp_dir),
+						kwargs={"queue": queue, "verbose": verbose}
+					)
+					process.start()
+					process.join()
+					return_values = queue.get(timeout=60)
+				else:
+					return_values = self._train_subprocess(dataset, checkpoint_path, temp_dir, verbose=verbose)
+				if return_values:
+					break
+			except Exception as e:
+				warning(f"Error during _train_fold: {e}\nRetrying in 60 seconds ({try_count + 1}/10)...")
+				time.sleep(60)
+		history: dict[str, Any] = return_values["history"]
+		eval_results: dict[str, Any] = return_values["eval_results"]
+		predictions: NDArray[Any] = return_values["predictions"]
+		true_classes: NDArray[Any] = return_values["true_classes"]
+		training_predictions: NDArray[Any] = return_values.get("training_predictions", None)
+		training_true_classes: NDArray[Any] = return_values.get("training_true_classes", None)
+
+		# For each epoch, log the history
+		mlflow_utils.log_history(history, prefix=mlflow_prefix)
+
+		# Append the history and evaluation results
+		self.history.append(history)
+		self.evaluation_results.append(eval_results)
+
+		# Generate and save ROC Curve and PR Curve for this fold
+		MetricUtils.all_curves(true_classes, predictions, fold_number, run_name=self.run_name)
+
+		# If final model, also log the ROC curve and PR curve for the train set
+		if fold_number == 0:
+			fold_number = -2	# -2 is the train set
+			MetricUtils.all_curves(training_true_classes, training_predictions, fold_number, run_name=self.run_name)
+
+		# Log visualization artifacts if they were generated
+		if temp_dir is not None:
+			mlflow.log_artifacts(temp_dir.name)
+			temp_dir.cleanup()
+
+		# Show some metrics
+		if verbose > 0:
+			last_history: dict[str, Any] = {k: v[-1] for k, v in history.items()}
+			info(f"Training done, metrics: {last_history}")
+
+		# Return the trained model
+		return return_values.get("model", None)
+