[ENH] Add Feature load_model to Deep Regressor Ensembles (#3130)

* fix: docs, fix: Correct misleading load_model docstring & resolve NameError in return type

* feat: implement load_model functionality to InceptionTimeRegressor and LiteTimeRegressor (Closes #2770)

* test: Add unit tests for load_model functionality (Issue #2770)

* style(docs): Resolve D205 ruff error by adding required blank line to all modified load_model docstrings.

* style(base): Final cleanup of base regressor docstring after merge.

* resolve return type in BaseDeepClassifier and BaseDeepRegressor

* revert return type of BaseDeepRegressor and BaseDeepClassifier to None

* Automatic `pre-commit` fixes

---------

Co-authored-by: rwtarpit <180088298+rwtarpit@users.noreply.github.com>

Author: rwtarpit

aeon/classification/deep_learning/_inception_time.py

@@ -1,5 +1,7 @@
 """InceptionTime and Inception classifiers."""
 
+from __future__ import annotations
+
 __maintainer__ = ["hadifawaz1999"]
 __all__ = ["InceptionTimeClassifier"]
 
@@ -251,7 +253,7 @@ def __init__(
 
         super().__init__()
 
-    def _fit(self, X, y):
+    def _fit(self, X: np.ndarray, y: np.ndarray) -> InceptionTimeClassifier:
         """Fit the ensemble of IndividualInceptionClassifier models.
 
         Parameters
@@ -307,7 +309,7 @@ def _fit(self, X, y):
 
         return self
 
-    def _predict(self, X) -> np.ndarray:
+    def _predict(self, X: np.ndarray) -> np.ndarray:
         """Predict the labels of the test set using InceptionTime.
 
         Parameters
@@ -328,7 +330,7 @@ def _predict(self, X) -> np.ndarray:
             ]
         )
 
-    def _predict_proba(self, X) -> np.ndarray:
+    def _predict_proba(self, X: np.ndarray) -> np.ndarray:
         """Predict the proba of labels of the test set using InceptionTime.
 
         Parameters
@@ -351,24 +353,30 @@ def _predict_proba(self, X) -> np.ndarray:
         return probs
 
     @classmethod
-    def load_model(self, model_path, classes):
-        """Load pre-trained classifiers instead of fitting.
+    def load_model(
+        self, model_path: list[str], classes: np.ndarray
+    ) -> InceptionTimeClassifier:
+        """Load pre-trained keras models from disk instead of fitting.
 
+        Pretrained models should be saved using "save_best_model"
+        or "save_last_model" boolean parameter.
         When calling this function, all functionalities can be used
-        such as predict, predict_proba, etc. with the loaded models.
+        such as predict, predict_proba etc. with the loaded model.
 
         Parameters
         ----------
         model_path : list of str (list of paths including the model names and extension)
-            The directory where the models will be saved including the model
-            names with a ".keras" extension.
-        classes : np.ndarray
+            The complete path (including file name and '.keras' extension)
+            from which the pre-trained model's weights and configuration
+            are loaded.
+         classes : np.ndarray
             The set of unique classes the pre-trained loaded model is trained
             to predict during the classification task.
+        Example: model_path="path/to/file/best_model.keras"
 
         Returns
         -------
-        None
+        InceptionTimeClassifier
         """
         assert (
             type(model_path) is list

aeon/classification/deep_learning/_lite_time.py

@@ -1,5 +1,7 @@
 """LITETime and LITE classifiers."""
 
+from __future__ import annotations
+
 __maintainer__ = ["hadifawaz1999"]
 __all__ = ["LITETimeClassifier"]
 
@@ -193,7 +195,7 @@ def __init__(
 
         super().__init__()
 
-    def _fit(self, X, y):
+    def _fit(self, X: np.ndarray, y: np.ndarray) -> LITETimeClassifier:
         """Fit the ensemble of IndividualLITEClassifier models.
 
         Parameters
@@ -239,7 +241,7 @@ def _fit(self, X, y):
 
         return self
 
-    def _predict(self, X) -> np.ndarray:
+    def _predict(self, X: np.ndarray) -> np.ndarray:
         """Predict the labels of the test set using LITETime.
 
         Parameters
@@ -260,7 +262,7 @@ def _predict(self, X) -> np.ndarray:
             ]
         )
 
-    def _predict_proba(self, X) -> np.ndarray:
+    def _predict_proba(self, X: np.ndarray) -> np.ndarray:
         """Predict the proba of labels of the test set using LITETime.
 
         Parameters
@@ -283,24 +285,30 @@ def _predict_proba(self, X) -> np.ndarray:
         return probs
 
     @classmethod
-    def load_model(self, model_path, classes):
-        """Load pre-trained classifiers instead of fitting.
+    def load_model(
+        self, model_path: list[str], classes: np.ndarray
+    ) -> LITETimeClassifier:
+        """Load pre-trained keras models from disk instead of fitting.
 
+        Pretrained models should be saved using "save_best_model"
+        or "save_last_model" boolean parameter.
         When calling this function, all functionalities can be used
-        such as predict, predict_proba, etc. with the loaded models.
+        such as predict, predict_proba etc. with the loaded model.
 
         Parameters
         ----------
         model_path : list of str (list of paths including the model names and extension)
-            The director where the models will be saved including the model
-            names with a ".keras" extension.
-        classes : np.ndarray
+            The complete path (including file name and '.keras' extension)
+            from which the pre-trained model's weights and configuration
+            are loaded.
+         classes : np.ndarray
             The set of unique classes the pre-trained loaded model is trained
             to predict during the classification task.
+        Example: model_path="path/to/file/best_model.keras"
 
         Returns
         -------
-        None
+        LITETimeClassifier
         """
         assert (
             type(model_path) is list

aeon/classification/deep_learning/base.py

@@ -22,6 +22,8 @@ class name: BaseDeepClassifier
 because we can generalise tags, _predict and _predict_proba
 """
 
+from __future__ import annotations
+
 __maintainer__ = ["hadifawaz1999"]
 __all__ = ["BaseDeepClassifier"]
 
@@ -178,17 +180,20 @@ def save_last_model_to_file(self, file_path="./"):
         """
         self.model_.save(file_path + self.last_file_name + ".keras")
 
-    def load_model(self, model_path, classes):
+    def load_model(self, model_path: str, classes: np.ndarray) -> None:
         """Load a pre-trained keras model instead of fitting.
 
+        Pretrained model should be saved using "save_last_model" or
+        "save_best_model" boolean parameter.
         When calling this function, all functionalities can be used
         such as predict, predict_proba etc. with the loaded model.
 
         Parameters
         ----------
         model_path : str (path including model name and extension)
-            The directory where the model will be saved including the model
-            name with a ".keras" extension.
+            The complete path (including file name and '.keras' extension)
+            from which the pre-trained model's weights and configuration
+            are loaded.
             Example: model_path="path/to/file/best_model.keras"
         classes : np.ndarray
             The set of unique classes the pre-trained loaded model is trained

aeon/regression/deep_learning/_inception_time.py

@@ -343,6 +343,43 @@ def _predict(self, X: np.ndarray) -> np.ndarray:
 
         return ypreds
 
+    @classmethod
+    def load_model(self, model_path: list[str]) -> InceptionTimeRegressor:
+        """Load pre-trained keras models from disk instead of fitting.
+
+        Pretrained models should be saved using "save_best_model"
+        or "save_last_model" boolean parameter.
+        When calling this function, all functionalities can be used
+        such as predict, etc. with the loaded model.
+
+        Parameters
+        ----------
+        model_path : list of str (list of paths including the model names and extension)
+            The complete path (including file name and '.keras' extension)
+            from which the pre-trained model's weights and configuration
+            are loaded.
+        Example: model_path="path/to/file/best_model.keras"
+
+        Returns
+        -------
+        InceptionTimeRegressor
+        """
+        assert (
+            type(model_path) is list
+        ), "model_path should be a list of paths to the models"
+
+        regressor = self()
+        regressor.regressors_ = []
+
+        for i in range(len(model_path)):
+            reg = IndividualInceptionRegressor()
+            reg.load_model(model_path[i])
+            regressor.regressors_.append(reg)
+
+        regressor.n_regressors = len(regressor.regressors_)
+        regressor.is_fitted = True
+        return regressor
+
     @classmethod
     def _get_test_params(
         cls, parameter_set: str = "default"

aeon/regression/deep_learning/_lite_time.py

@@ -269,6 +269,43 @@ def _predict(self, X: np.ndarray) -> np.ndarray:
 
         return vals
 
+    @classmethod
+    def load_model(self, model_path: list[str]) -> LITETimeRegressor:
+        """Load pre-trained keras models from disk instead of fitting.
+
+        Pretrained models should be saved using "save_best_model"
+        or "save_last_model" boolean parameter.
+        When calling this function, all functionalities can be used
+        such as predict, etc. with the loaded model.
+
+        Parameters
+        ----------
+        model_path : list of str (list of paths including the model names and extension)
+            The complete path (including file name and '.keras' extension)
+            from which the pre-trained model's weights and configuration
+            are loaded.
+        Example: model_path="path/to/file/best_model.keras"
+
+        Returns
+        -------
+        LITETimeRegressor
+        """
+        assert (
+            type(model_path) is list
+        ), "model_path should be a list of paths to the models"
+
+        regressor = self()
+        regressor.regressors_ = []
+
+        for i in range(len(model_path)):
+            reg = IndividualLITERegressor()
+            reg.load_model(model_path[i])
+            regressor.regressors_.append(reg)
+
+        regressor.n_regressors = len(regressor.regressors_)
+        regressor.is_fitted = True
+        return regressor
+
     @classmethod
     def _get_test_params(cls, parameter_set: str = "default") -> dict | list[dict]:
         """Return testing parameter settings for the estimator.

aeon/regression/deep_learning/base.py

@@ -120,14 +120,17 @@ def save_last_model_to_file(self, file_path: str = "./") -> None:
     def load_model(self, model_path: str) -> None:
         """Load a pre-trained keras model instead of fitting.
 
+        Pretrained model should be saved using "save_last_model"
+        or "save_best_model" boolean parameter.
         When calling this function, all functionalities can be used
-        such as predict etc. with the loaded model.
+        such as predict, etc. with the loaded model.
 
         Parameters
         ----------
         model_path : str (path including model name and extension)
-            The directory where the model will be saved including the model
-            name with a ".keras" extension.
+            The complete path (including file name and '.keras' extension)
+            from which the pre-trained model's weights and configuration
+            are loaded.
             Example: model_path="path/to/file/best_model.keras"
 
         Returns

aeon/regression/deep_learning/tests/test_inception_time.py

@@ -0,0 +1,49 @@
+"""Tests for save/load functionality of InceptionTimeRegressor."""
+
+import glob
+import os
+import tempfile
+
+import numpy as np
+import pytest
+
+from aeon.regression.deep_learning import InceptionTimeRegressor
+from aeon.testing.data_generation import make_example_3d_numpy
+from aeon.utils.validation._dependencies import _check_soft_dependencies
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("tensorflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_save_load_inceptiontime():
+    """Test saving and loading for InceptionTimeRegressor."""
+    with tempfile.TemporaryDirectory() as temp:
+        temp_dir = os.path.join(temp, "")
+
+        X, y = make_example_3d_numpy(
+            n_cases=10,
+            n_channels=1,
+            n_timepoints=12,
+            return_y=True,
+            regression_target=True,
+        )
+
+        model = InceptionTimeRegressor(
+            n_epochs=1, random_state=42, save_best_model=True, file_path=temp_dir
+        )
+        model.fit(X, y)
+
+        y_pred_orig = model.predict(X)
+
+        model_file = glob.glob(os.path.join(temp_dir, f"{model.best_file_name}*.keras"))
+
+        loaded_model = InceptionTimeRegressor.load_model(model_path=model_file)
+
+        assert isinstance(loaded_model, InceptionTimeRegressor)
+
+        preds = loaded_model.predict(X)
+        assert isinstance(preds, np.ndarray)
+
+        assert len(preds) == len(y)
+        np.testing.assert_array_equal(preds, y_pred_orig)

aeon/regression/deep_learning/tests/test_lite_time.py

@@ -0,0 +1,51 @@
+"""Tests for save/load functionality of LiteTimeRegressor."""
+
+import glob
+import os
+import tempfile
+
+import numpy as np
+import pytest
+
+from aeon.regression.deep_learning import LITETimeRegressor
+from aeon.testing.data_generation import make_example_3d_numpy
+from aeon.utils.validation._dependencies import _check_soft_dependencies
+
+
+@pytest.mark.skipif(
+    not _check_soft_dependencies("tensorflow", severity="none"),
+    reason="skip test if required soft dependency not available",
+)
+def test_save_load_litetim():
+    """Test saving and loading for LiteTimeRegressor."""
+    with tempfile.TemporaryDirectory() as temp:
+        temp_dir = os.path.join(temp, "")
+
+        X, y = make_example_3d_numpy(
+            n_cases=10,
+            n_channels=1,
+            n_timepoints=12,
+            return_y=True,
+            regression_target=True,
+        )
+
+        model = LITETimeRegressor(
+            n_epochs=1, random_state=42, save_best_model=True, file_path=temp_dir
+        )
+        model.fit(X, y)
+
+        y_pred_orig = model.predict(X)
+
+        model_files = glob.glob(
+            os.path.join(temp_dir, f"{model.best_file_name}*.keras")
+        )
+
+        loaded_model = LITETimeRegressor.load_model(model_path=model_files)
+
+        assert isinstance(loaded_model, LITETimeRegressor)
+
+        preds = loaded_model.predict(X)
+        assert isinstance(preds, np.ndarray)
+
+        assert len(preds) == len(y)
+        np.testing.assert_array_equal(preds, y_pred_orig)