Repository: zurutech/ashpy

Issue:
The model selection JSON file is overwritten with defaults on re-start
Scenario:

- You're training a model and doing model selection
- You stop the train
- You re-start the train

The JSON file of the model selection is overwritten with default values, making you lost the previously (correctly stored) value.

Discussion:

Patch:
diff --git a/src/ashpy/contexts/classifier.py b/src/ashpy/contexts/classifier.py
index 649bb33..27588ca 100644
--- a/src/ashpy/contexts/classifier.py
+++ b/src/ashpy/contexts/classifier.py
@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import tensorflow as tf  # pylint: disable=import-error
 from ashpy.contexts.context import Context
@@ -36,7 +36,7 @@ class ClassifierContext(Context):
         classifier_model: tf.keras.Model = None,
         loss: ClassifierLoss = None,  # ?: Do we really need to default these values to None?
         dataset: tf.data.Dataset = None,
-        metrics: List[Metric] = None,
+        metrics: Tuple[Metric] = None,
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
         global_step: tf.Variable = tf.Variable(
             0, name="global_step", trainable=False, dtype=tf.int64
diff --git a/src/ashpy/contexts/context.py b/src/ashpy/contexts/context.py
index 08d03a1..31f56be 100644
--- a/src/ashpy/contexts/context.py
+++ b/src/ashpy/contexts/context.py
@@ -20,7 +20,7 @@ collections of variable encapsulated in a Python Class as a way to seamlessly
 handle information transfer.
 """
 
-from typing import List, Optional
+from typing import Optional, Tuple
 
 import tensorflow as tf
 from ashpy.metrics import Metric
@@ -32,7 +32,7 @@ class Context:
 
     def __init__(
         self,
-        metrics: List[Metric] = None,
+        metrics: Tuple[Metric] = None,
         dataset: tf.data.Dataset = None,
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
         global_step=tf.Variable(0, name="global_step", trainable=False, dtype=tf.int64),
@@ -42,7 +42,7 @@ class Context:
         Initialize the Context.
 
         Args:
-            metrics (:obj:`list` of [:py:class:`ashpy.metrics.metric.Metric`]): List of
+            metrics (:obj:`tuple` of (:py:class:`ashpy.metrics.metric.Metric`)): List of
                 :py:class:`ashpy.metrics.metric.Metric` objects.
             dataset (:py:class:`tf.data.Dataset`): The dataset to use, that
                 contains everything needed to use the model in this context.
@@ -55,8 +55,7 @@ class Context:
         """
         self._distribute_strategy = tf.distribute.get_strategy()
 
-        # TODO: are metrics really needed right now?
-        self._metrics = metrics if metrics else []
+        self._metrics = metrics if metrics else ()
         self._dataset = dataset
         self._log_eval_mode = log_eval_mode
         self._global_step = global_step
@@ -98,12 +97,12 @@ class Context:
         self._dataset = _dataset
 
     @property
-    def metrics(self) -> List[Metric]:
+    def metrics(self) -> Tuple[Metric]:
         """
         Retrieve the metrics.
 
         Returns:
-            :obj:`list` of [:py:class:`ashpy.metrics.metric.Metric`].
+            :obj:`tuple` of (:py:class:`ashpy.metrics.metric.Metric`).
 
         """
         return self._metrics
diff --git a/src/ashpy/metrics/metric.py b/src/ashpy/metrics/metric.py
index 296e59b..41273e4 100644
--- a/src/ashpy/metrics/metric.py
+++ b/src/ashpy/metrics/metric.py
@@ -113,13 +113,13 @@ class Metric(ABC):
         # write the initial value of the best metric
         if not self.best_model_sel_file.exists():
             self.best_model_sel_file.parent.mkdir(parents=True)
-        initial_value = (
-            np.inf if self._model_selection_operator is operator.lt else -np.inf
-        )
-        self.json_write(
-            self.best_model_sel_file,
-            {self.sanitized_name: str(initial_value), "step": 0},
-        )
+            initial_value = (
+                np.inf if self._model_selection_operator is operator.lt else -np.inf
+            )
+            self.json_write(
+                self.best_model_sel_file,
+                {self.sanitized_name: str(initial_value), "step": 0},
+            )
 
     @property
     def name(self) -> str:
diff --git a/src/ashpy/restorers/__init__.py b/src/ashpy/restorers/__init__.py
index 6aed16f..e7f02b2 100644
--- a/src/ashpy/restorers/__init__.py
+++ b/src/ashpy/restorers/__init__.py
@@ -45,11 +45,12 @@ Restorers allow for easy restoration of tracked objects from :class:`tf.train.Ch
 
 from ashpy.restorers.classifier import ClassifierRestorer
 from ashpy.restorers.gan import AdversarialEncoderRestorer, AdversarialRestorer
-from ashpy.restorers.restorer import Restorer
+from ashpy.restorers.restorer import ModelNotConstructedError, Restorer
 
 __ALL__ = [
     "Restorer",
     "AdversarialRestorer",
     "AdversarialEncoderRestorer",
     "ClassifierRestorer",
+    "ModelNotConstructedError",
 ]
diff --git a/src/ashpy/restorers/restorer.py b/src/ashpy/restorers/restorer.py
index 518dff2..339c44e 100644
--- a/src/ashpy/restorers/restorer.py
+++ b/src/ashpy/restorers/restorer.py
@@ -21,7 +21,22 @@ from typing import Dict, List, Optional, Union
 import ashpy
 import tensorflow as tf
 
-__ALL__ = ["Restorer"]
+__ALL__ = ["Restorer, ModelNotConstructedError"]
+
+
+class ModelNotConstructedError(Exception):
+    """
+    Exception raised while restoring sub-classed Model before having called it on data.
+
+    Warning:
+        When restoring a :class:`tf.keras.Model` object from checkpoint assure that the
+        model has been correctly built and instantiated by firstly calling it on some
+        sample inputs. In the case of a model built with either the Sequential or
+        Functional API an exception will be raised; for a model built with the Chainer API
+        it will fail silently, restoration will be "successful" but no values will actually
+        be restored since there are no valid placeholder as the model has not be built yet.
+
+    """
 
 
 class Restorer:
@@ -95,6 +110,27 @@ class Restorer:
                 f"Object {placeholder} is should be of type: {placeholder_type}"
             )
 
+    @staticmethod
+    def _check_model_construction(restored_model: tf.keras.Model) -> bool:
+        """
+        Optimistically check that the model.weights property returns a non empty-list.
+
+        The underlying assumption is that Models created via the sub-classing API, when restored
+        without being properly constructed AKA called on some input, will have empty lists
+        as layers.weights.
+
+        TODO: add docs for the exception.
+        TODO: add test case for the Sequential without input shape
+        """
+        try:
+            if restored_model.weights == []:
+                raise ModelNotConstructedError
+        except AttributeError:
+            # A Sequential() buil without specifiyng the input shape can be treated as a
+            # sub-classed model for restoration purposes.
+            raise ModelNotConstructedError
+        return True
+
     def restore_object(self, placeholder, object_ckpt_id: str):
         """
         Restore a placeholder from a checkpoint using the specified id.
@@ -112,6 +148,8 @@ class Restorer:
         """
         checkpoint = tf.train.Checkpoint(**{object_ckpt_id: placeholder})
         status = self._restore_checkpoint(checkpoint)
+        if isinstance(placeholder, tf.keras.Model):
+            assert self._check_model_construction(placeholder)
         print(self._restored_log_msg.format(object_ckpt_id, self._ckpts_dir))
         return status
 
diff --git a/src/ashpy/trainers/classifier.py b/src/ashpy/trainers/classifier.py
index 8f2b2c7..3932624 100644
--- a/src/ashpy/trainers/classifier.py
+++ b/src/ashpy/trainers/classifier.py
@@ -15,7 +15,7 @@
 """Primitive Trainer Interface."""
 
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
 import ashpy
 import tensorflow as tf
@@ -41,7 +41,7 @@ class ClassifierTrainer(Trainer):
         optimizer: tf.optimizers.Optimizer,
         loss: ashpy.losses.ClassifierLoss,
         epochs: int,
-        metrics: Optional[List[Metric]] = None,
+        metrics: Optional[Union[Tuple[Metric], List[Metric]]] = None,
         callbacks: Optional[List[Callback]] = None,
         logdir: Union[Path, str] = Path().cwd() / "log",
         global_step: Optional[tf.Variable] = None,
@@ -56,7 +56,7 @@ class ClassifierTrainer(Trainer):
             loss (:obj:`ashpy.losses.classifier.ClassifierLoss`): A loss function built following
                 :py:mod:`ashpy.executors``.
             epochs (int): Number of training epochs.
-            metrics: (List): List of :py:class:`ashpy.metrics.metric.Metric` to
+            metrics: (Tuple/List): Tuple/List of :py:class:`ashpy.metrics.metric.Metric` to
                 measure on training and validation data.
             callbacks (List): List of :py:class:`ashpy.callbacks.callback.Callback` to
                 to call on events
@@ -136,9 +136,9 @@ class ClassifierTrainer(Trainer):
 
         self._avg_loss = ClassifierLoss(name="ashpy/avg_loss")
         if metrics:
-            metrics.append(self._avg_loss)
+            metrics = (*metrics, self._avg_loss)
         else:
-            metrics = [self._avg_loss]
+            metrics = (self._avg_loss,)
 
         super()._update_metrics(metrics)
         super()._validate_metrics()
@@ -160,6 +160,14 @@ class ClassifierTrainer(Trainer):
             checkpoint=self._checkpoint,
         )
 
+    def _build_and_restore_models(self, dataset: tf.data.Dataset):
+        restorer = ashpy.restorers.ClassifierRestorer(self._logdir)
+        (x, _) = next(iter(dataset.take(1)))
+        # Invoke model on sample input
+        self._model(x)
+        restorer.restore_model(self._model)
+        self._deferred_restoration = False
+
     def train_step(self, features, labels):
         """
         Train step.
@@ -212,6 +220,9 @@ class ClassifierTrainer(Trainer):
                 performance.
 
         """
+        if self._deferred_restoration:
+            self._build_and_restore_models(dataset=training_set)
+
         # set the context properties
         self._context.training_set = training_set
         self._context.validation_set = validation_set
diff --git a/src/ashpy/trainers/gan.py b/src/ashpy/trainers/gan.py
index 09e6e2e..8098128 100644
--- a/src/ashpy/trainers/gan.py
+++ b/src/ashpy/trainers/gan.py
@@ -14,8 +14,9 @@
 
 """Collection of GANs trainers."""
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import List, Optional, Tuple, Union
 
+import ashpy.restorers
 import tensorflow as tf
 from ashpy.callbacks import Callback
 from ashpy.contexts.gan import GANContext, GANEncoderContext
@@ -138,7 +139,7 @@ class AdversarialTrainer(Trainer):
         generator_loss: Executor,
         discriminator_loss: Executor,
         epochs: int,
-        metrics: Optional[List[Metric]] = None,
+        metrics: Optional[Union[Tuple[Metric], List[Metric]]] = None,
         callbacks: Optional[List[Callback]] = None,
         logdir: Union[Path, str] = Path().cwd() / "log",
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
@@ -191,12 +192,12 @@ class AdversarialTrainer(Trainer):
         self._discriminator_loss = discriminator_loss
         self._discriminator_loss.reduction = tf.losses.Reduction.NONE
 
-        losses_metrics = [
+        losses_metrics = (
             DiscriminatorLoss(name="ashpy/d_loss", logdir=logdir),
             GeneratorLoss(name="ashpy/g_loss", logdir=logdir),
-        ]
+        )
         if metrics:
-            metrics.extend(losses_metrics)
+            metrics = (*metrics, *losses_metrics)
         else:
             metrics = losses_metrics
 
@@ -229,6 +230,16 @@ class AdversarialTrainer(Trainer):
             metrics=self._metrics,
         )
 
+    def _build_and_restore_models(self, dataset: tf.data.Dataset):
+        restorer = ashpy.restorers.AdversarialRestorer(self._logdir)
+        (x, _), z = next(iter(dataset.take(1)))
+        # Invoke model on sample input
+        self._generator(z)
+        self._discriminator(x)
+        restorer.restore_generator(self._generator)
+        restorer.restore_discriminator(self._discriminator)
+        self._deferred_restoration = False
+
     def train_step(self, real_xy, g_inputs):
         """
         Train step for the AdversarialTrainer.
@@ -317,6 +328,9 @@ class AdversarialTrainer(Trainer):
                     performance.
 
         """
+        if self._deferred_restoration:
+            self._build_and_restore_models(dataset=dataset)
+
         current_epoch = self._current_epoch()
 
         self._update_global_batch_size(
@@ -575,6 +589,17 @@ class EncoderTrainer(AdversarialTrainer):
             metrics=self._metrics,
         )
 
+    def _build_and_restore_models(self, dataset: tf.data.Dataset):
+        restorer = ashpy.restorers.AdversarialEncoderRestorer(self._logdir)
+        (x, _), _ = next(iter(dataset.take(1)))
+
+        # Invoke model on sample input
+        self._encoder(x)
+        restorer.restore_encoder(self._encoder)
+
+        super()._build_and_restore_models(dataset)
+        self._deferred_restoration = False
+
     def train_step(self, real_xy, g_inputs):
         """Adversarial training step.
 
@@ -668,6 +693,9 @@ class EncoderTrainer(AdversarialTrainer):
                 performance.
 
         """
+        if self._deferred_restoration:
+            self._build_and_restore_models(dataset=dataset)
+
         current_epoch = self._current_epoch()
 
         self._update_global_batch_size(
diff --git a/src/ashpy/trainers/trainer.py b/src/ashpy/trainers/trainer.py
index 3522247..b08748d 100644
--- a/src/ashpy/trainers/trainer.py
+++ b/src/ashpy/trainers/trainer.py
@@ -19,6 +19,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
+import ashpy.restorers
 import tensorflow as tf
 from ashpy.callbacks import Callback
 from ashpy.contexts import Context
@@ -36,6 +37,7 @@ class Trainer(ABC):
     ckpt_id_global_step: str = "global_step"
     ckpt_id_steps_per_epoch: str = "steps_per_epoch"
     ckpt_id_callbacks: str = "callbacks"
+    _deferred_restoration: bool = False
 
     def __init__(
         self,
@@ -44,8 +46,8 @@ class Trainer(ABC):
         logdir: Union[Path, str] = Path().cwd() / "log",
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
         global_step: Optional[tf.Variable] = None,
-        metrics: Optional[List[Metric]] = None,
-        callbacks: Optional[List[Callback]] = None,
+        metrics: Optional[Union[Tuple[Metric], List[Metric]]] = None,
+        callbacks: Optional[Union[Tuple[Callback], List[Callback]]] = None,
     ) -> None:
         r"""
         Primitive trainer interface. Handles model saving and restore.
@@ -62,9 +64,10 @@ class Trainer(ABC):
                 to use when evaluating and logging.
             global_step (Optional[py:class:`ashpy.modes.LogEvalMode`]): tf.Variable that
                 keeps track of the training steps.
-            metrics (Optional[List[:py:class:`ashpy.metrics.Metric`]]): list of metrics.
-            callbacks (Optional[List[:py:class:`ashpy.callbacks.Callback`]]): list of callbacks
-                to handle events.
+            metrics (Optional[List | Tuple [:py:class:`ashpy.metrics.Metric`]]): list or
+                tuple of metrics.
+            callbacks (Optional[List[:py:class:`ashpy.callbacks.Callback`]]): list or
+                tuple of callbacks to handle events.
 
         """
         self._distribute_strategy = tf.distribute.get_strategy()
@@ -72,14 +75,14 @@ class Trainer(ABC):
 
         # set and validate metrics
         if metrics is None:
-            metrics = []
+            metrics = ()
         self._metrics = metrics
         self._validate_metrics()
 
         # set and validate callbacks
         if callbacks is None:
-            callbacks = []
-        self._callbacks: List[Callback] = callbacks
+            callbacks = ()
+        self._callbacks: Tuple[Callback] = tuple(callbacks)
         self._validate_callbacks()
 
         self._epochs = epochs
@@ -245,7 +248,23 @@ class Trainer(ABC):
     def _restore_or_init(self):
         """Restore or initialize the persistence layer (checkpoint)."""
         if self._manager.latest_checkpoint:
-            self._checkpoint.restore(self._manager.latest_checkpoint)
+            restorer = ashpy.restorers.Restorer(self._logdir)
+
+            # Check that at least the ids of the things we want to restore are present in
+            # the map
+            for k in self._ckpt_dict:
+                assert k in restorer.checkpoint_map
+                try:
+                    restorer.restore_object(self._ckpt_dict[k], k)
+                except ashpy.restorers.ModelNotConstructedError:
+                    print(
+                        "WARNING: Trying to restore a Model constructed via sub-classing"
+                        "API requires calling it on some data first."
+                        "More info -> "
+                        "https://www.tensorflow.org/guide/keras/save_and_serialize#saving_subclassed_models"  # pylint: disable=line-too-long
+                        "Restoration will be deferred at the beginning of the Trainer call."
+                    )
+                    self._deferred_restoration = True
             print(f"Restored checkpoint {self._manager.latest_checkpoint}.")
         else:
             print("Initializing checkpoint.")
@@ -372,6 +391,10 @@ class Trainer(ABC):
                 )
         return tuple(columns)
 
+    @abstractmethod
+    def _build_and_restore_models(self, dataset: tf.data.Dataset):
+        """Build and restore a Subclassed model by firstly calling it on some data."""
+
     @abstractmethod
     def call(self, *args, **kwargs):
         """Execute the training process."""


Test patch:
diff --git a/conftest.py b/conftest.py
index d12e717..5c5580a 100644
--- a/conftest.py
+++ b/conftest.py
@@ -30,8 +30,8 @@ from ashpy.metrics import (
     SSIM_Multiscale,
 )
 from tests.utils.fake_training_loop import (
-    fake_adversarial_training_loop,
-    fake_classifier_training_loop,
+    FakeAdversarialTraining,
+    FakeClassifierTraining,
 )
 
 
@@ -70,7 +70,7 @@ def save_dir():
 TEST_MATRIX = {
     # NOTE: Always pass metrics as Tuple, Trainers produce side effects!
     "adversarial_trainer": [
-        fake_adversarial_training_loop,
+        FakeAdversarialTraining,
         {
             "image_resolution": [256, 256],
             "layer_spec_input_res": (8, 8),
@@ -101,7 +101,7 @@ TEST_MATRIX = {
         ),
     ],
     "classifier_trainer": [
-        fake_classifier_training_loop,
+        FakeClassifierTraining,
         {"measure_performance_freq": 1},
         (ClassifierLoss(model_selection_operator=operator.lt),),
     ],
@@ -112,8 +112,11 @@ LOOPS = [TEST_MATRIX[k] for k in TEST_MATRIX]
 
 
 @pytest.fixture(scope="function", params=LOOPS, ids=TRAINING_IDS)
-def fake_training(request):
+def fake_training_fn(request):
     """Fixture used to generate fake training for the tests."""
     training_loop, loop_args, metrics = request.param
     assert len(metrics) in [1, 3]
-    return (training_loop, loop_args, list(metrics))
+
+    return lambda logdir, **kwargs: training_loop(
+        logdir=logdir, metrics=metrics, **loop_args, **kwargs
+    )
diff --git a/tests/callbacks/__init__.py b/tests/callbacks/__init__.py
index 67de119..d76e7f9 100644
--- a/tests/callbacks/__init__.py
+++ b/tests/callbacks/__init__.py
@@ -1,1 +1,6 @@
-"""Test for callbacks."""
+"""
+Test for callbacks.
+
+TODO: test a callback that retrieves a Metric from its Context
+TODO: Restructure these tests
+"""
diff --git a/tests/callbacks/test_callbacks.py b/tests/callbacks/test_callbacks.py
index b382511..4b28a08 100644
--- a/tests/callbacks/test_callbacks.py
+++ b/tests/callbacks/test_callbacks.py
@@ -17,10 +17,10 @@
 from ashpy.callbacks import LogImageGANCallback
 from ashpy.callbacks.events import Event
 
-from tests.utils.fake_training_loop import fake_adversarial_training_loop
+from tests.utils.fake_training_loop import FakeAdversarialTraining
 
 
 def test_callbacks(tmpdir):
     """Test the integration between callbacks and trainer."""
     callbacks = [LogImageGANCallback(event=Event.ON_BATCH_END, event_freq=1)]
-    fake_adversarial_training_loop(tmpdir, callbacks=callbacks)
+    FakeAdversarialTraining(tmpdir, callbacks=callbacks)
diff --git a/tests/callbacks/test_counter_callback.py b/tests/callbacks/test_counter_callback.py
index 03f4d41..57e48f0 100644
--- a/tests/callbacks/test_counter_callback.py
+++ b/tests/callbacks/test_counter_callback.py
@@ -18,7 +18,7 @@ import pytest
 from ashpy.callbacks import CounterCallback, Event
 from ashpy.models.gans import ConvDiscriminator, ConvGenerator
 
-from tests.utils.fake_training_loop import fake_adversarial_training_loop
+from tests.utils.fake_training_loop import FakeAdversarialTraining
 
 
 class FakeCounterCallback(CounterCallback):
@@ -81,11 +81,11 @@ def test_counter_callback(_models, tmpdir):
     )
     callbacks = [clbk]
     generator, discriminator = _models
-    fake_adversarial_training_loop(
+    FakeAdversarialTraining(
         logdir=tmpdir,
         callbacks=callbacks,
         generator=generator,
         discriminator=discriminator,
         epochs=1,
-    )
+    )()
     assert clbk.fake_counter == 1
diff --git a/tests/callbacks/test_custom_callback.py b/tests/callbacks/test_custom_callback.py
index ede8e99..f838fb8 100644
--- a/tests/callbacks/test_custom_callback.py
+++ b/tests/callbacks/test_custom_callback.py
@@ -18,7 +18,7 @@ import pytest
 from ashpy.callbacks import Callback
 from ashpy.callbacks.events import Event
 
-from tests.utils.fake_training_loop import fake_adversarial_training_loop
+from tests.utils.fake_training_loop import FakeAdversarialTraining
 
 
 class MCallback(Callback):
@@ -65,13 +65,13 @@ def test_custom_callbacks(tmpdir, event: Event):
     dataset_size = 2
     batch_size = 2
 
-    fake_adversarial_training_loop(
+    FakeAdversarialTraining(
         logdir=tmpdir,
         callbacks=callbacks,
         epochs=epochs,
         dataset_size=dataset_size,
         batch_size=batch_size,
-    )
+    )()
 
     # assert the number of times the on_event has been called
     assert m_callback.counter == get_n_events_from_epochs(
diff --git a/tests/callbacks/test_save_callback.py b/tests/callbacks/test_save_callback.py
index f450609..82b1d5b 100644
--- a/tests/callbacks/test_save_callback.py
+++ b/tests/callbacks/test_save_callback.py
@@ -21,7 +21,7 @@ import pytest
 from ashpy.callbacks import SaveCallback, SaveFormat, SaveSubFormat
 from ashpy.models.gans import ConvDiscriminator, ConvGenerator
 
-from tests.utils.fake_training_loop import fake_adversarial_training_loop
+from tests.utils.fake_training_loop import FakeAdversarialTraining
 
 COMPATIBLE_FORMAT_AND_SUB_FORMAT = [
     (SaveFormat.WEIGHTS, SaveSubFormat.TF),
@@ -106,9 +106,9 @@ def _test_save_callback_helper(tmpdir, save_format, save_sub_format, save_dir: P
         )
     ]
 
-    fake_adversarial_training_loop(
+    FakeAdversarialTraining(
         tmpdir, callbacks=callbacks, generator=generator, discriminator=discriminator,
-    )
+    )()
 
 
 def test_save_callback_type_error(save_dir: str,):
diff --git a/tests/test_losses.py b/tests/test_losses.py
index 01d316a..9fbf13c 100644
--- a/tests/test_losses.py
+++ b/tests/test_losses.py
@@ -21,7 +21,7 @@ from ashpy.losses.gan import (
     get_adversarial_loss_generator,
 )
 
-from tests.utils.fake_training_loop import fake_adversarial_training_loop
+from tests.utils.fake_training_loop import FakeAdversarialTraining
 
 
 @pytest.mark.parametrize("loss_type", list(AdversarialLossType))
@@ -31,6 +31,6 @@ def test_losses(loss_type: AdversarialLossType, tmpdir):
     generator_loss = get_adversarial_loss_generator(loss_type)()
     discriminator_loss = get_adversarial_loss_discriminator(loss_type)()
 
-    fake_adversarial_training_loop(
+    FakeAdversarialTraining(
         tmpdir, generator_loss=generator_loss, discriminator_loss=discriminator_loss,
-    )
+    )()
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index 93540a3..dfa2f26 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -19,13 +19,13 @@ TODO: Adversarial Encoder Traner
 """
 import json
 import operator
-import pathlib
 import shutil
+from pathlib import Path
 
 import pytest
 from ashpy.metrics import ClassifierLoss
 
-from tests.utils.fake_training_loop import fake_classifier_training_loop
+from tests.utils.fake_training_loop import FakeAdversarialTraining, FakeTraining
 
 DEFAULT_LOGDIR = "log"
 OPERATOR_INITIAL_VALUE_MAP = {operator.gt: "-inf", operator.lt: "inf"}
@@ -34,7 +34,7 @@ OPERATOR_INITIAL_VALUE_MAP = {operator.gt: "-inf", operator.lt: "inf"}
 @pytest.fixture(scope="module")
 def cleanup():
     """Remove the default logdir before and after testing."""
-    default_log_dir = pathlib.Path(DEFAULT_LOGDIR)
+    default_log_dir = Path(DEFAULT_LOGDIR)
     if default_log_dir.exists():
         shutil.rmtree(default_log_dir)
     yield "Cleanup"
@@ -42,7 +42,22 @@ def cleanup():
         shutil.rmtree(default_log_dir)
 
 
-def test_metrics_log(fake_training, tmpdir, cleanup):
+def get_metric_data(metric, tmpdir):
+    """Make sure that the metric exists and is readable. Return its value."""
+    metric_dir = Path(tmpdir) / "best" / metric.sanitized_name
+    assert metric_dir.exists()
+    json_path = metric_dir / f"{metric.sanitized_name}.json"
+    assert json_path.exists()
+    with open(json_path, "r") as fp:
+        metric_data = json.load(fp)
+
+        # Assert the metric data contains the expected keys
+        assert metric.sanitized_name in metric_data
+        assert "step" in metric_data
+        return metric_data[metric.sanitized_name]
+
+
+def test_metrics_log(fake_training_fn, tmpdir, cleanup):
     """
     Test that trainers correctly create metrics log files.
 
@@ -57,52 +72,74 @@ def test_metrics_log(fake_training, tmpdir, cleanup):
         THEN the values of the keys should not be the operator initial value
 
     """
-    training_loop, loop_args, metrics = fake_training
+    logdir = Path(tmpdir)
 
-    training_completed, trainer = training_loop(
-        logdir=tmpdir, metrics=metrics, **loop_args
-    )
+    fake_training: FakeTraining = fake_training_fn(logdir=logdir)
+    assert fake_training()
+    trainer = fake_training.trainer
 
-    assert training_completed
-    assert not pathlib.Path(DEFAULT_LOGDIR).exists()  # Assert absence of side effects
+    assert not Path(DEFAULT_LOGDIR).exists()  # Assert absence of side effects
     # Assert there exists folder for each metric
     for metric in trainer._metrics:
-        metric_dir = pathlib.Path(tmpdir) / "best" / metric.sanitized_name
-        assert metric_dir.exists()
-        json_path = metric_dir / f"{metric.sanitized_name}.json"
-        assert json_path.exists()
-        with open(json_path, "r") as fp:
-            metric_data = json.load(fp)
-
-            # Assert the metric data contains the expected keys
-            assert metric.sanitized_name in metric_data
-            assert "step" in metric_data
-
-            # Assert that the correct model selection has been performed
-            # Check it by seeing if the values in the json has been updated
-            if metric.model_selection_operator:
-                try:
-                    initial_value = OPERATOR_INITIAL_VALUE_MAP[
-                        metric.model_selection_operator
-                    ]
-                except KeyError:
-                    raise ValueError(
-                        "Please add the initial value for this operator to OPERATOR_INITIAL_VALUE_MAP"
-                    )
-                assert metric_data[metric.sanitized_name] != initial_value
+
+        metric_value = get_metric_data(metric, tmpdir)
+        # Assert that the correct model selection has been performed
+        # Check it by seeing if the values in the json has been updated
+        if metric.model_selection_operator:
+            try:
+                initial_value = OPERATOR_INITIAL_VALUE_MAP[
+                    metric.model_selection_operator
+                ]
+            except KeyError:
+                raise ValueError(
+                    "Please add the initial value for this "
+                    "operator to OPERATOR_INITIAL_VALUE_MAP"
+                )
+
+            assert metric_value != initial_value
 
 
 # -------------------------------------------------------------------------------------
 
 
-@pytest.mark.parametrize("training_loop", [fake_classifier_training_loop])
-def test_metrics_names_collision(training_loop, tmpdir):
+def test_metrics_names_collision(tmpdir):
     """
     Test that an exception is correctly raised when two metrics have the same name.
 
     WHEN two or more metrics passed to a trainer have the same name
         THEN raise a ValueError
     """
-    metrics = [ClassifierLoss(name="test_loss"), ClassifierLoss(name="test_loss")]
+    metrics = [
+        ClassifierLoss(name="test_loss"),
+        ClassifierLoss(name="test_loss"),
+    ]
+
     with pytest.raises(ValueError):
-        training_loop(metrics=metrics, logdir=tmpdir)
+        # Since names collision is checked by the primitive Trainer we can test it easily
+        # with just one FakeTraining and be assured that it works.
+        FakeAdversarialTraining(tmpdir, metrics=metrics)
+
+
+# -------------------------------------------------------------------------------------
+
+
+def test_metrics_on_restart(fake_training_fn, tmpdir):
+    """Test that metrics are correctly read on train restart."""
+
+    fake_training = fake_training_fn(tmpdir)
+    assert fake_training()
+
+    t1_values = {
+        metric.name: get_metric_data(metric, tmpdir)
+        for metric in fake_training.trainer._metrics
+    }
+    print(t1_values)
+
+    restart = fake_training_fn(tmpdir)
+
+    t2_values = {
+        metric.name: get_metric_data(metric, tmpdir)
+        for metric in restart.trainer._metrics
+    }
+    print(t2_values)
+    assert t1_values == t2_values
diff --git a/tests/test_restorers.py b/tests/test_restorers.py
index c7851a4..51c91bb 100644
--- a/tests/test_restorers.py
+++ b/tests/test_restorers.py
@@ -22,38 +22,36 @@ GIVEN some training has been done
     WHEN restoring models the first layer of the restored model and the trained one should
         have the same weights.
 """
-
 from pathlib import Path
+from typing import Union
 
 import pytest
 import tensorflow as tf
 from ashpy.callbacks import CounterCallback
-from ashpy.restorers import AdversarialRestorer, ClassifierRestorer, Restorer
+from ashpy.restorers import (
+    AdversarialRestorer,
+    ClassifierRestorer,
+    ModelNotConstructedError,
+    Restorer,
+)
 from ashpy.trainers import AdversarialTrainer, ClassifierTrainer
 
-from tests.utils.fake_datasets import (
-    fake_adversarial_dataset,
-    fake_autoencoder_datasest,
+from tests.utils.fake_training_loop import (
+    FakeAdversarialTraining,
+    FakeClassifierTraining,
 )
-from tests.utils.fake_models import basic_dcgan, conv_autoencoder
 
 DEFAULT_CKPT_DIR = "ckpts"
 
 
-def _check_first_layer(trained: tf.keras.Model, restored: tf.keras.Model, i=0):
+def _check_models_weights(trained: tf.keras.Model, restored: tf.keras.Model, i=0):
     """Test that the first layers of the restored and trained model have the same weights."""
+
     try:
-        try:
-            trained_layer = trained.layers[i].weights[0]
-            restored_layers = restored.layers[i].weights[0]
-        except IndexError:
-            trained_layer = trained.layers[i].weights
-            restored_layers = restored.layers[i].weightss
-        assert tf.reduce_all(tf.equal(trained_layer, restored_layers))
-    except AttributeError:
-        i += 1
-        print(f"Proceed to layer in position {i}")
-        _check_first_layer(trained, restored, i)
+        for i, element in enumerate(trained.weights):
+            assert tf.reduce_all(tf.equal(element, restored.weights[i]))
+    except AssertionError:
+        raise ModelNotConstructedError
 
 
 def _test_restore_object(restorer, placeholder, ckpt_id, capsys):
@@ -71,50 +69,62 @@ def _check_log(restorer, ckpt_id, capsys):
     )
 
 
-def test_restore_model(fake_training, capsys, tmpdir):
+def test_restore_model(fake_training_fn, capsys, tmpdir):
     """
     Test that models are correctly restored.
 
     The test is performed by checking the logs and the first layer of each model.
     """
-    training_loop, loop_args, metrics = fake_training
-    training_completed, trainer = training_loop(
-        logdir=tmpdir, metrics=metrics, **loop_args
-    )
-    assert training_completed
-    restorer = Restorer(logdir=tmpdir)
+    logdir = Path(tmpdir).joinpath("training")
+    _tmp_logdir = Path(tmpdir).joinpath("banana")
+
+    fake_training = fake_training_fn(logdir=logdir)
+    assert fake_training()
+    trainer = fake_training.trainer
+    restorer = Restorer(logdir=logdir)
 
     if isinstance(trainer, ClassifierTrainer):
 
-        placeholder = conv_autoencoder()
+        new_loop = fake_training_fn(logdir=_tmp_logdir)
+        placeholder = new_loop.model
 
         # Ensure model have been built correctly
-        x, _ = next(iter(fake_autoencoder_datasest()))
+        x, _ = next(iter(new_loop.dataset))
         placeholder(x)
 
         ckpt_id = trainer.ckpt_id_model
         _test_restore_object(restorer, placeholder, ckpt_id, capsys)
-        _check_first_layer(trainer._model, placeholder)
+        _check_models_weights(trainer._model, placeholder)
 
     elif isinstance(trainer, AdversarialTrainer):
 
-        placeholder_g, placeholder_d = basic_dcgan(**loop_args)
+        new_loop: FakeAdversarialTraining = fake_training_fn(logdir=_tmp_logdir)
+        placeholder_g, placeholder_d = (new_loop.generator, new_loop.discriminator)
+
+        # Ensure that the ModelNotConstructedError is correctly triggered
+        with pytest.raises(ModelNotConstructedError):
+            _test_restore_object(
+                restorer, placeholder_g, trainer.ckpt_id_generator, capsys
+            )
+            _test_restore_object(
+                restorer, placeholder_d, trainer.ckpt_id_discriminator, capsys
+            )
 
         # Ensure model have been built correctly
-        (x, _), z = next(iter(fake_adversarial_dataset(**loop_args)))
+        (x, _), z = next(iter(new_loop.dataset))
         fake = placeholder_g(z)
         assert tf.reduce_all(tf.equal(fake.shape, x.shape))
         placeholder_d(x)
 
         _test_restore_object(restorer, placeholder_g, trainer.ckpt_id_generator, capsys)
-        _check_first_layer(trainer._generator, placeholder_g)
+        _check_models_weights(trainer._generator, placeholder_g)
         _test_restore_object(
             restorer, placeholder_d, trainer.ckpt_id_discriminator, capsys
         )
-        _check_first_layer(trainer._discriminator, placeholder_d)
+        _check_models_weights(trainer._discriminator, placeholder_d)
 
 
-def test_restore_common_variables(fake_training, capsys, tmpdir):
+def test_restore_common_variables(fake_training_fn, capsys, tmpdir):
     """
     Test that the convenience methods exposed by :class:`Restorer` work correctly.
 
@@ -123,12 +133,12 @@ def test_restore_common_variables(fake_training, capsys, tmpdir):
         - steps per epoch
 
     """
-    training_loop, loop_args, metrics = fake_training
-    training_completed, trainer = training_loop(
-        logdir=tmpdir, metrics=metrics, **loop_args
-    )
-    assert training_completed
-    restorer = Restorer(logdir=tmpdir)
+    logdir = Path(tmpdir).joinpath("training")
+
+    fake_training = fake_training_fn(logdir=logdir)
+    assert fake_training()
+    trainer = fake_training.trainer
+    restorer = Restorer(logdir=logdir)
 
     # Restore variables and check their values using the convenience method
     assert tf.equal(trainer._global_step, restorer.get_global_step())
@@ -147,17 +157,17 @@ def test_restore_common_variables(fake_training, capsys, tmpdir):
         ) in out.split("\n")
 
 
-def test_restore_callbacks(fake_training, capsys, tmpdir):
+def test_restore_callbacks(fake_training_fn, capsys, tmpdir):
     """Test that callbacks are succesfully restored."""
-    training_loop, loop_args, metrics = fake_training
-    training_completed, trainer = training_loop(
-        logdir=tmpdir, metrics=metrics, **loop_args
-    )
-    assert training_completed
-    restorer = Restorer(logdir=tmpdir)
+    logdir = Path(tmpdir).joinpath("training")
+
+    fake_training = fake_training_fn(logdir=logdir)
+    assert fake_training()
+    trainer = fake_training.trainer
+    restorer = Restorer(logdir=logdir)
 
     if isinstance(trainer, AdversarialTrainer):
-        placeholder_callbacks = loop_args.get("callbacks")
+        placeholder_callbacks = fake_training.callbacks
         for i, placeholder_callback in enumerate(placeholder_callbacks):
             # Restore the callbacks
             restorer.restore_callback(placeholder_callback, placeholder_callback.name)
@@ -171,14 +181,14 @@ def test_restore_callbacks(fake_training, capsys, tmpdir):
                 )
 
 
-def test_read_checkpoint_map(fake_training, tmpdir):
+def test_read_checkpoint_map(fake_training_fn, tmpdir):
     """Test that checkpoint map is read correctly."""
-    training_loop, loop_args, metrics = fake_training
-    training_completed, trainer = training_loop(
-        logdir=tmpdir, metrics=metrics, **loop_args
-    )
-    assert training_completed
-    restorer = Restorer(logdir=tmpdir)
+    logdir = Path(tmpdir).joinpath("training")
+
+    fake_training = fake_training_fn(logdir=logdir)
+    assert fake_training()
+    trainer = fake_training.trainer
+    restorer = Restorer(logdir=logdir)
     assert restorer.checkpoint_map == trainer._generate_checkpoint_map()
 
     # Test that Restorer.checkpoint_map without the checkpoint_map.json correctly returns None
@@ -186,7 +196,7 @@ def test_read_checkpoint_map(fake_training, tmpdir):
     ckpt_map: Path = restorer._ckpts_dir / "checkpoint_map.json"
     ckpt_map.unlink()
     assert not ckpt_map.exists()
-    assert not Restorer(tmpdir).checkpoint_map
+    assert not Restorer(logdir).checkpoint_map
 
 
 # ###################################################
@@ -203,7 +213,7 @@ def _test_convenience_model_restorer(
 ):
     convenience_method(placeholder_model)
     _check_log(restorer, ckpt_id, capsys)
-    _check_first_layer(trained_model, placeholder_model)
+    _check_models_weights(trained_model, placeholder_model)
 
 
 def _test_convenience_optimizer_restorer(
@@ -218,27 +228,30 @@ def _test_convenience_optimizer_restorer(
     _check_log(restorer, ckpt_id, capsys)
 
 
-def test_convenience_restorer(fake_training, capsys, tmpdir):
+def test_convenience_restorer(fake_training_fn, capsys, tmpdir):
     """
     Test that models and optimizers are correctly restored using the convenience classes.
 
     TODO: Add test for AdversarialEncoderRestorer
     """
-    logdir = tmpdir
-    training_loop, loop_args, metrics = fake_training
-    training_completed, trainer = training_loop(
-        logdir=logdir, metrics=metrics, **loop_args
-    )
-    assert training_completed
+    logdir = Path(tmpdir).joinpath("training")
+    _tmp_logdir = Path(tmpdir).joinpath("banana")
+
+    fake_training = fake_training_fn(logdir=logdir)
+    assert fake_training()
+    trainer = fake_training.trainer
+    restorer = Restorer(logdir=logdir)
 
     if isinstance(trainer, ClassifierTrainer):
         restorer: ClassifierRestorer = ClassifierRestorer(logdir=logdir)
 
-        placeholder_model = conv_autoencoder()
+        new_training: FakeClassifierTraining = fake_training_fn(_tmp_logdir)
+
+        placeholder_model = new_training.model
         placeholder_opt = tf.keras.optimizers.Adam()
 
         # Ensure model have been built correctly
-        x, _ = next(iter(fake_autoencoder_datasest()))
+        x, _ = next(iter(new_training.dataset))
         placeholder_model(x)
 
         _test_convenience_model_restorer(
@@ -260,14 +273,40 @@ def test_convenience_restorer(fake_training, capsys, tmpdir):
     elif isinstance(trainer, AdversarialTrainer):
         restorer: AdversarialRestorer = AdversarialRestorer(logdir=logdir)
 
-        placeholder_g, placeholder_d = basic_dcgan(**loop_args)
+        new_training: FakeAdversarialTraining = fake_training_fn(_tmp_logdir)
+
+        placeholder_g, placeholder_d = (
+            new_training.generator,
+            new_training.discriminator,
+        )
         placeholder_optimizer_g, placeholder_optimizer_d = (
             tf.keras.optimizers.Adam(),
             tf.keras.optimizers.Adam(),
         )
 
-        # Ensure model have been built correctly
-        (x, _), z = next(iter(fake_adversarial_dataset(**loop_args)))
+        # Ensure that the ModelNotConstructedError is correctly triggered
+        with pytest.raises(ModelNotConstructedError):
+            _test_convenience_model_restorer(
+                restorer,
+                restorer.restore_generator,
+                placeholder_g,
+                trainer._generator,
+                trainer.ckpt_id_generator,
+                capsys,
+            )
+
+        with pytest.raises(ModelNotConstructedError):
+            _test_convenience_model_restorer(
+                restorer,
+                restorer.restore_discriminator,
+                placeholder_d,
+                trainer._discriminator,
+                trainer.ckpt_id_discriminator,
+                capsys,
+            )
+
+        # Ensure models have been built correctly
+        (x, _), z = next(iter(new_training.dataset))
         fake = placeholder_g(z)
         assert tf.reduce_all(tf.equal(fake.shape, x.shape))
         placeholder_d(x)
diff --git a/tests/test_trainers.py b/tests/test_trainers.py
index d0dd6cb..5a43432 100644
--- a/tests/test_trainers.py
+++ b/tests/test_trainers.py
@@ -18,20 +18,59 @@ from pathlib import Path
 from typing import List
 
 import ashpy
+import pytest
+import tensorflow as tf
+from ashpy.trainers import AdversarialTrainer, ClassifierTrainer
 
+from tests.test_restorers import ModelNotConstructedError, _check_models_weights
+from tests.utils.fake_training_loop import (
+    FakeAdversarialTraining,
+    FakeClassifierTraining,
+    FakeTraining,
+)
 
-def test_generate_human_ckpt_dict(fake_training, tmpdir):
+
+def test_correct_trainer_restoration_on_restart(fake_training_fn, tmpdir):
+    logdir = Path(tmpdir)
+
+    fake_training: FakeTraining = fake_training_fn(logdir=logdir)
+    assert fake_training()
+
+    if isinstance(fake_training, FakeClassifierTraining):
+        assert isinstance(fake_training.trainer, ClassifierTrainer)
+        trained_model = fake_training.trainer._model
+
+        new_training: FakeClassifierTraining = fake_training_fn(logdir=logdir)
+        new_training.trainer._build_and_restore_models(new_training.dataset)
+        restored_model = new_training.trainer._model
+
+        _check_models_weights(trained_model, restored_model)
+
+    if isinstance(fake_training, FakeAdversarialTraining):
+        assert isinstance(fake_training.trainer, AdversarialTrainer)
+        trained_g = fake_training.trainer._generator
+        trained_d = fake_training.trainer._discriminator
+
+        new_training: FakeAdversarialTraining = fake_training_fn(logdir=logdir)
+        new_training.trainer._build_and_restore_models(new_training.dataset)
+        restored_g = new_training.trainer._generator
+        restored_d = new_training.trainer._discriminator
+        _check_models_weights(trained_g, restored_g)
+        _check_models_weights(trained_d, restored_d)
+
+
+def test_generate_human_ckpt_dict(fake_training_fn, tmpdir):
     """
     Test that the generation of the human readable map of the ckpt_dict works.
 
     TODO: improve the test.
     """
     logdir = Path(tmpdir)
-    training_loop, loop_args, metrics = fake_training
-    training_completed, trainer = training_loop(
-        logdir=logdir, metrics=metrics, **loop_args
-    )
-    assert training_completed
+    fake_training = fake_training_fn(logdir=logdir)
+    assert fake_training()
+
+    trainer = fake_training.trainer
+
     assert trainer._checkpoint_map
     assert (Path(trainer._ckpts_dir) / "checkpoint_map.json").exists()
     metrics: List[ashpy.metrics.Metric] = trainer._metrics
diff --git a/tests/utils/fake_training_loop.py b/tests/utils/fake_training_loop.py
index 74ce155..f749056 100644
--- a/tests/utils/fake_training_loop.py
+++ b/tests/utils/fake_training_loop.py
@@ -15,13 +15,13 @@
 """Fake training loop to simplify training in tests."""
 import operator
 from pathlib import Path
-from typing import Union
+from typing import List, Tuple, Union
 
 import ashpy
 import tensorflow as tf
 from ashpy.losses import DiscriminatorMinMax, GeneratorBCE
 from ashpy.models.gans import ConvDiscriminator, ConvGenerator
-from ashpy.trainers import AdversarialTrainer
+from ashpy.trainers import AdversarialTrainer, ClassifierTrainer, Trainer
 
 from tests.utils.fake_datasets import (
     fake_adversarial_dataset,
@@ -29,131 +29,164 @@ from tests.utils.fake_datasets import (
 )
 from tests.utils.fake_models import basic_dcgan, conv_autoencoder
 
-__ALL__ = ["fake_classifier_training_loop", "fake_adversarial_training_loop"]
-
-
-def fake_classifier_training_loop(
-    # Trainer
-    logdir: Union[Path, str] = "testlog",
-    optimizer=tf.optimizers.Adam(1e-4),
-    metrics=[ashpy.metrics.ClassifierLoss(model_selection_operator=operator.lt)],
-    epochs=2,
-    # Dataset
-    dataset_size=10,
-    image_resolution=(64, 64),
-    batch_size=5,
-    # Model: Autoencoder
-    layer_spec_input_res=(64, 64),
-    layer_spec_target_res=(4, 4),
-    kernel_size=3,
-    initial_filters=16,
-    filters_cap=64,
-    encoding_dimension=50,
-    channels=3,
-    # Call parameters
-    measure_performance_freq=10,
-):
-    """Fake Classifier training loop implementation using an autoencoder as a base model."""
-    # Model
-    model = conv_autoencoder(
-        layer_spec_input_res,
-        layer_spec_target_res,
-        kernel_size,
-        initial_filters,
-        filters_cap,
-        encoding_dimension,
-        channels,
-    )
-
-    # Dataset
-    dataset = fake_autoencoder_datasest(
-        dataset_size, image_resolution, channels, batch_size
-    )
-
-    # Loss
-    reconstruction_error = ashpy.losses.ClassifierLoss(
-        tf.keras.losses.MeanSquaredError()
-    )
-
-    # Trainer
-    trainer = ashpy.trainers.ClassifierTrainer(
-        model=model,
-        optimizer=optimizer,
-        loss=reconstruction_error,
-        logdir=str(logdir),
-        epochs=epochs,
-        metrics=metrics,
-    )
-
-    trainer(dataset, dataset, measure_performance_freq=measure_performance_freq)
-    return 1, trainer
-
-
-def fake_adversarial_training_loop(
-    logdir: Union[Path, str] = "testlog",
-    kernel_size=(5, 5),
-    metrics=None,
-    callbacks=None,
-    epochs=2,
-    dataset_size=2,
-    batch_size=2,
-    generator_loss=GeneratorBCE(),
-    discriminator_loss=DiscriminatorMinMax(),
-    image_resolution=(28, 28),
-    layer_spec_input_res=(7, 7),
-    layer_spec_target_res=(7, 7),
-    channels=1,
-    output_shape=1,
-    latent_dim=100,
-    # Call parameters
-    measure_performance_freq=10,
-    # Models from outside
-    generator=None,
-    discriminator=None,
-):
-    """Fake training loop implementation."""
-    # test parameters
-    if callbacks is None:
-        callbacks = []
-    if metrics is None:
-        metrics = []
-
-    # Model definition
-    models = basic_dcgan(
-        image_resolution=image_resolution,
-        layer_spec_input_res=layer_spec_input_res,
-        layer_spec_target_res=layer_spec_target_res,
-        kernel_size=kernel_size,
-        channels=channels,
-        output_shape=output_shape,
-    )
-    if not generator:
-        generator = models[0]
-    if not discriminator:
-        discriminator = models[1]
-
-    # Trainer
-    trainer = AdversarialTrainer(
-        generator=generator,
-        discriminator=discriminator,
-        generator_optimizer=tf.optimizers.Adam(1e-4),
-        discriminator_optimizer=tf.optimizers.Adam(1e-4),
-        generator_loss=generator_loss,
-        discriminator_loss=discriminator_loss,
-        epochs=epochs,
-        metrics=metrics,
-        callbacks=callbacks,
-        logdir=logdir,
-    )
-
-    dataset = fake_adversarial_dataset(
-        image_resolution=image_resolution,
-        epochs=epochs,
-        dataset_size=dataset_size,
-        batch_size=batch_size,
-        latent_dim=latent_dim,
-        channels=channels,
-    )
-
-    trainer(dataset, measure_performance_freq=measure_performance_freq)
-    return 1, trainer
+__ALL__ = ["FakeTraining", "FakeAdversarialTraining", "FakeClassifierTraining"]
+
+
+class FakeTraining:
+    dataset: tf.data.Dataset
+    trainer: Trainer
+    measure_performance_freq: int
+    metrics: Union[List[ashpy.metrics.Metric], Tuple[ashpy.metrics.Metric]]
+
+
+class FakeClassifierTraining(FakeTraining):
+    def __init__(
+        self,
+        # Trainer
+        logdir: Union[Path, str] = "testlog",
+        optimizer=tf.optimizers.Adam(1e-4),
+        metrics=[ashpy.metrics.ClassifierLoss(model_selection_operator=operator.lt)],
+        epochs=2,
+        # Dataset
+        dataset_size=10,
+        image_resolution=(64, 64),
+        batch_size=5,
+        # Model: Autoencoder
+        layer_spec_input_res=(64, 64),
+        layer_spec_target_res=(4, 4),
+        kernel_size=3,
+        initial_filters=16,
+        filters_cap=64,
+        encoding_dimension=50,
+        channels=3,
+        # Call parameters
+        measure_performance_freq=10,
+    ):
+        """Fake Classifier training loop implementation using an autoencoder as a base model."""
+        self.measure_performance_freq = measure_performance_freq
+
+        # Model
+        self.model: tf.keras.Model = conv_autoencoder(
+            layer_spec_input_res,
+            layer_spec_target_res,
+            kernel_size,
+            initial_filters,
+            filters_cap,
+            encoding_dimension,
+            channels,
+        )
+
+        # Dataset
+        self.dataset = fake_autoencoder_datasest(
+            dataset_size, image_resolution, channels, batch_size
+        )
+
+        # Loss
+        reconstruction_error = ashpy.losses.ClassifierLoss(
+            tf.keras.losses.MeanSquaredError()
+        )
+
+        # Trainer
+        self.trainer = ClassifierTrainer(
+            model=self.model,
+            optimizer=optimizer,
+            loss=reconstruction_error,
+            logdir=str(logdir),
+            epochs=epochs,
+            metrics=metrics,
+        )
+        self.metrics = metrics
+
+    def __call__(self) -> bool:
+        self.trainer(
+            self.dataset,
+            self.dataset,
+            measure_performance_freq=self.measure_performance_freq,
+        )
+        return True
+
+
+# ---------------------------------------------------------------------------------
+
+
+class FakeAdversarialTraining(FakeTraining):
+    def __init__(
+        self,
+        logdir: Union[Path, str] = "testlog",
+        kernel_size=(5, 5),
+        metrics=None,
+        callbacks=None,
+        epochs=2,
+        dataset_size=2,
+        batch_size=2,
+        generator_loss=GeneratorBCE(),
+        discriminator_loss=DiscriminatorMinMax(),
+        image_resolution=(28, 28),
+        layer_spec_input_res=(7, 7),
+        layer_spec_target_res=(7, 7),
+        channels=1,
+        output_shape=1,
+        latent_dim=100,
+        # Call parameters
+        measure_performance_freq=10,
+        # Models from outside
+        generator=None,
+        discriminator=None,
+    ):
+        """Fake training loop implementation."""
+        self.measure_performance_freq = measure_performance_freq
+
+        # test parameters
+        if callbacks is None:
+            callbacks = []
+        if metrics is None:
+            metrics = []
+
+        self.metrics = metrics
+        self.callbacks = callbacks
+
+        # Model definition
+        models = basic_dcgan(
+            image_resolution=image_resolution,
+            layer_spec_input_res=layer_spec_input_res,
+            layer_spec_target_res=layer_spec_target_res,
+            kernel_size=kernel_size,
+            channels=channels,
+            output_shape=output_shape,
+        )
+        if not generator:
+            generator = models[0]
+        if not discriminator:
+            discriminator = models[1]
+
+        self.generator = generator
+        self.discriminator = discriminator
+
+        # Trainer
+        self.trainer = AdversarialTrainer(
+            generator=generator,
+            discriminator=discriminator,
+            generator_optimizer=tf.optimizers.Adam(1e-4),
+            discriminator_optimizer=tf.optimizers.Adam(1e-4),
+            generator_loss=generator_loss,
+            discriminator_loss=discriminator_loss,
+            epochs=epochs,
+            metrics=metrics,
+            callbacks=callbacks,
+            logdir=logdir,
+        )
+
+        self.dataset = fake_adversarial_dataset(
+            image_resolution=image_resolution,
+            epochs=epochs,
+            dataset_size=dataset_size,
+            batch_size=batch_size,
+            latent_dim=latent_dim,
+            channels=channels,
+        )
+
+    def __call__(self) -> bool:
+        self.trainer(
+            self.dataset, measure_performance_freq=self.measure_performance_freq
+        )
+        return True