Repository: zurutech/ashpy

Issue:
No such file or directory checkpoint_map.json when restoring model from model selection
The title says pretty much everything.

Creating a restorer (in particular a ClassifierRestorer but it is not important) passing as checkpoint directory, the directory where the checkpoint of the best model has been saved causes an error.

two options:

1. We have to create the same file also when doing model selection
2. We have to make it possible to restore the model from the checkpoint even if the file is not present

I like the first one

Discussion:

Patch:
diff --git a/examples/gans/pix2pix_facades.py b/examples/gans/pix2pix_facades.py
index 5254933..b9f00be 100644
--- a/examples/gans/pix2pix_facades.py
+++ b/examples/gans/pix2pix_facades.py
@@ -17,7 +17,7 @@ Pix2Pix on Facades Datasets dummy implementation.
 
 Input Pipeline taken from: https://www.tensorflow.org/beta/tutorials/generative/pix2pix
 """
-import os
+from pathlib import Path
 
 import tensorflow as tf
 
@@ -34,7 +34,7 @@ from ashpy.trainers.gan import AdversarialTrainer
 _URL = "https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz"
 
 PATH_TO_ZIP = tf.keras.utils.get_file("facades.tar.gz", origin=_URL, extract=True)
-PATH = os.path.join(os.path.dirname(PATH_TO_ZIP), "facades/")
+PATH = Path(PATH_TO_ZIP).parent / "facades"
 
 BUFFER_SIZE = 100
 BATCH_SIZE = 1
@@ -172,10 +172,10 @@ def main(
     )
 
     metrics = []
-    logdir = f'{"log"}/{dataset_name}/run2'
+    logdir = Path("log") / dataset_name / "run2"
 
-    if not os.path.exists(logdir):
-        os.makedirs(logdir)
+    if not logdir.exists():
+        logdir.mkdir(parents=True)
 
     trainer = AdversarialTrainer(
         generator=generator,
diff --git a/examples/gans/pix2pix_facades_multi_gpu.py b/examples/gans/pix2pix_facades_multi_gpu.py
index 5492ee9..c8833a9 100644
--- a/examples/gans/pix2pix_facades_multi_gpu.py
+++ b/examples/gans/pix2pix_facades_multi_gpu.py
@@ -96,8 +96,8 @@ def main(
         metrics = []
         logdir = f'{"log"}/{dataset_name}/run_multi'
 
-        if not os.path.exists(logdir):
-            os.makedirs(logdir)
+        if not logdir.exists():
+            logdir.mkdir(parents=True)
 
         trainer = AdversarialTrainer(
             generator=generator,
diff --git a/src/ashpy/callbacks/gan.py b/src/ashpy/callbacks/gan.py
index 1636c82..45fb748 100644
--- a/src/ashpy/callbacks/gan.py
+++ b/src/ashpy/callbacks/gan.py
@@ -39,7 +39,7 @@ class LogImageGANCallback(CounterCallback):
 
             import shutil
             import operator
-            import os
+            from pathlib import Path
 
             generator = models.gans.ConvGenerator(
                 layer_spec_input_res=(7, 7),
@@ -69,7 +69,7 @@ class LogImageGANCallback(CounterCallback):
 
             # Trainer
             epochs = 2
-            logdir = "testlog/callbacks"
+            logdir = Path("testlog/callbacks")
             callbacks = [callbacks.LogImageGANCallback()]
             trainer = trainers.gan.AdversarialTrainer(
                 generator=generator,
@@ -100,7 +100,7 @@ class LogImageGANCallback(CounterCallback):
             trainer(dataset)
             shutil.rmtree(logdir)
 
-            assert not os.path.exists(logdir)
+            assert not logdir.exists()
 
             trainer._global_step.assign_add(500)
 
diff --git a/src/ashpy/callbacks/save_callback.py b/src/ashpy/callbacks/save_callback.py
index 5f8e844..3e6958e 100644
--- a/src/ashpy/callbacks/save_callback.py
+++ b/src/ashpy/callbacks/save_callback.py
@@ -13,10 +13,10 @@
 # limitations under the License.
 
 """Save weights callback."""
-import os
 import shutil
 from collections import deque
 from enum import Enum, Flag, auto
+from pathlib import Path
 from typing import List
 
 import tensorflow as tf
@@ -47,23 +47,24 @@ class SaveFormat(Flag):
         return "saved-model-and-weights"
 
     @staticmethod
-    def _initialize_dirs(save_dir, save_format, save_sub_format):
+    def _initialize_dirs(save_dir: Path, save_format, save_sub_format) -> Path:
         """Initialize the directory for this save_format and sub-format."""
-        save_dir = os.path.join(save_dir, save_format.name())
-        if not os.path.exists(save_dir):
-            os.makedirs(save_dir)
+        save_dir = save_dir / save_format.name()
+        if not save_dir.exists():
+            save_dir.mkdir(parents=True)
 
         save_dir = (
             save_dir
             if save_sub_format == SaveSubFormat.TF
-            else os.path.join(save_dir, save_format.name())
+            else save_dir / save_format.name()
         )
+
         return save_dir
 
     def save(
         self,
         model: tf.keras.models.Model,
-        save_dir: str,
+        save_dir: Path,
         save_sub_format: SaveSubFormat = SaveSubFormat.TF,
     ) -> None:
         """
@@ -71,7 +72,7 @@ class SaveFormat(Flag):
 
         Args:
             model (:py:class:`tf.keras.models.Model`): model to Save.
-            save_dir (str): path of the file in which to save the model.
+            save_dir (:class:`pathlib.Path`): path of the file in which to save the model.
             save_sub_format (:py:class:`ashpy.callbacks.save_callback.SaveSubFormat`): sub-format
                 of the save operation.
 
@@ -81,14 +82,18 @@ class SaveFormat(Flag):
             save_dir = self._initialize_dirs(
                 save_dir, SaveFormat.WEIGHTS, save_sub_format
             )
-            model.save_weights(save_dir, save_format=save_sub_format.value)
+            # NOTE: Keras (TF 2.1.0) checks for h5 file using endswith attribute.
+            # Explicit conversion to strings is required
+            model.save_weights(str(save_dir), save_format=save_sub_format.value)
 
         if SaveFormat.MODEL & self:
 
             save_dir = self._initialize_dirs(
                 save_dir, SaveFormat.MODEL, save_sub_format
             )
-            model.save(save_dir, save_format=save_sub_format.value)
+            # NOTE: TensorFlow 2.1.0 wanth either binary or unicod string.
+            # Explicit conversion to strings is required
+            model.save(str(save_dir), save_format=save_sub_format.value)
 
         if not (SaveFormat.MODEL & self) | (SaveFormat.WEIGHTS & self):
             raise NotImplementedError(
@@ -138,7 +143,7 @@ class SaveCallback(CounterCallback):
 
     def __init__(
         self,
-        save_dir: str,
+        save_dir: Path,
         models: List[tf.keras.models.Model],
         event: Event = Event.ON_EPOCH_END,
         event_freq: int = 1,
@@ -248,10 +253,10 @@ class SaveCallback(CounterCallback):
                 )
 
             # Create the correct directory name
-            save_dir_i = os.path.join(self._save_dir, f"model-{i}-step-{step}")
+            save_dir_i = self._save_dir / f"model-{i}-step-{step}"
 
-            if not os.path.exists(save_dir_i):
-                os.makedirs(save_dir_i)
+            if not save_dir_i.exists():
+                save_dir_i.mkdir(parents=True)
 
             # Add to the history
             self._save_path_histories[i].append(save_dir_i)
diff --git a/src/ashpy/metrics/classifier.py b/src/ashpy/metrics/classifier.py
index 972207d..c3609da 100644
--- a/src/ashpy/metrics/classifier.py
+++ b/src/ashpy/metrics/classifier.py
@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-import os
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, Union
 
 import tensorflow as tf  # pylint: disable=import-error
@@ -28,6 +28,8 @@ if TYPE_CHECKING:
 
     TPRocessingPredictions = Dict[str, Union[Callable, Dict[str, Any]]]
 
+__ALL__ = ["ClassifierLoss", "ClassifierMetric"]
+
 
 class ClassifierLoss(Metric):
     """A handy way to measure the classification loss."""
@@ -36,7 +38,7 @@ class ClassifierLoss(Metric):
         self,
         name: str = "loss",
         model_selection_operator: Callable = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
     ) -> None:
         """
         Initialize the Metric.
@@ -89,7 +91,7 @@ class ClassifierMetric(Metric):
         self,
         metric: tf.keras.metrics.Metric,
         model_selection_operator: Callable = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         processing_predictions=None,
     ) -> None:
         """
diff --git a/src/ashpy/metrics/gan.py b/src/ashpy/metrics/gan.py
index c12190c..02eb008 100644
--- a/src/ashpy/metrics/gan.py
+++ b/src/ashpy/metrics/gan.py
@@ -18,7 +18,8 @@ from __future__ import annotations
 import operator
 import os
 import types
-from typing import TYPE_CHECKING, Callable
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Union
 
 import tensorflow as tf
 import tensorflow_hub as hub
@@ -31,6 +32,14 @@ if TYPE_CHECKING:
         GANEncoderContext,
     )
 
+__ALL__ = [
+    "DiscriminatorLoss",
+    "EncoderLoss",
+    "EncodingAccuracy",
+    "GeneratorLoss",
+    "InceptionScore",
+]
+
 
 class DiscriminatorLoss(Metric):
     """The Discriminator loss value."""
@@ -39,7 +48,7 @@ class DiscriminatorLoss(Metric):
         self,
         name: str = "d_loss",
         model_selection_operator: Callable = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
     ) -> None:
         """
         Initialize the Metric.
@@ -104,7 +113,7 @@ class GeneratorLoss(Metric):
         self,
         name: str = "g_loss",
         model_selection_operator: Callable = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
     ):
         """
         Initialize the Metric.
@@ -168,7 +177,7 @@ class EncoderLoss(Metric):
         self,
         name: str = "e_loss",
         model_selection_operator: Callable = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
     ) -> None:
         """
         Initialize the Metric.
@@ -241,7 +250,7 @@ class InceptionScore(Metric):
         inception: tf.keras.Model,
         name: str = "inception_score",
         model_selection_operator=operator.gt,
-        logdir=os.path.join(os.getcwd(), "log"),
+        logdir=Path().cwd() / "log",
     ):
         """
         Initialize the Metric.
@@ -355,7 +364,7 @@ class InceptionScore(Metric):
             from_logits=True
         ),
         optimizer: tf.keras.optimizers.Adam = tf.keras.optimizers.Adam(1e-5),
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
     ) -> tf.keras.Model:
         """
         Restore or train (and save) the Inception model.
@@ -396,7 +405,7 @@ class InceptionScore(Metric):
         ckpt.objects.extend([model, step])
         logdir = logdir
         manager = tf.train.CheckpointManager(
-            ckpt, os.path.join(logdir, "inception", name), max_to_keep=1
+            ckpt, logdir / "inception", name, max_to_keep=1
         )
 
         if manager.latest_checkpoint:
@@ -428,7 +437,7 @@ class EncodingAccuracy(ClassifierMetric):
         classifier: tf.keras.Model,
         name: str = "encoding_accuracy",
         model_selection_operator: Callable = None,
-        logdir=os.path.join(os.getcwd(), "log"),
+        logdir=Path().cwd() / "log",
     ) -> None:
         """
         Measure the Generator and Encoder performance together.
diff --git a/src/ashpy/metrics/metric.py b/src/ashpy/metrics/metric.py
index 5d0d8eb..296e59b 100644
--- a/src/ashpy/metrics/metric.py
+++ b/src/ashpy/metrics/metric.py
@@ -21,6 +21,7 @@ import json
 import operator
 import os
 from abc import ABC, abstractmethod
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
 import numpy as np
@@ -29,6 +30,8 @@ import tensorflow as tf  # pylint: disable=import-error
 if TYPE_CHECKING:
     from ashpy.contexts import Context
 
+__ALL__ = ["Metric"]
+
 
 class Metric(ABC):
     """
@@ -42,7 +45,7 @@ class Metric(ABC):
         name: str,
         metric: tf.keras.metrics.Metric,
         model_selection_operator: Callable = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path.cwd() / "log",
     ) -> None:
         """
         Initialize the Metric object.
@@ -66,11 +69,11 @@ class Metric(ABC):
         self._name = name
         self._metric = metric
         self._model_selection_operator = model_selection_operator
-        self._logdir = logdir
+        self._logdir = Path(logdir) if not isinstance(logdir, Path) else logdir
 
     def model_selection(
         self, checkpoint: tf.train.Checkpoint, global_step: tf.Variable
-    ) -> None:
+    ) -> Optional[Path]:
         """
         Perform model selection.
 
@@ -99,16 +102,17 @@ class Metric(ABC):
                 },
             )
             manager = tf.train.CheckpointManager(
-                checkpoint, os.path.join(self.best_folder, "ckpts"), max_to_keep=1
+                checkpoint, self.best_folder / "ckpts", max_to_keep=1
             )
-            manager.save()
+            return Path(manager.save())
+        return None
 
     def _update_logdir(self):
         if not self._model_selection_operator:
             pass
         # write the initial value of the best metric
-        if not os.path.exists(self.best_model_sel_file):
-            os.makedirs(os.path.dirname(self.best_model_sel_file))
+        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
         )
@@ -143,7 +147,7 @@ class Metric(ABC):
         return self._model_selection_operator
 
     @property
-    def logdir(self) -> str:
+    def logdir(self) -> Path:
         """Retrieve the log directory."""
         return self._logdir
 
@@ -154,17 +158,17 @@ class Metric(ABC):
         self._update_logdir()
 
     @property
-    def best_folder(self) -> str:
+    def best_folder(self) -> Path:
         """Retrieve the folder used to save the best model when doing model selection."""
-        return os.path.join(self.logdir, "best", self.sanitized_name)
+        return self.logdir / "best" / self.sanitized_name
 
     @property
-    def best_model_sel_file(self) -> str:
+    def best_model_sel_file(self) -> Path:
         """Retrieve the path to JSON file containing the measured performance of the best model."""
-        return os.path.join(self.best_folder, self.sanitized_name + ".json")
+        return self.best_folder / (self.sanitized_name + ".json")
 
     @staticmethod
-    def json_read(filename: str) -> Dict[str, Any]:
+    def json_read(filename: Path) -> Dict[str, Any]:
         """
         Read a JSON file.
 
@@ -175,7 +179,7 @@ class Metric(ABC):
             :py:obj:`typing.Dict`: Dictionary containing the content of the JSON file.
 
         """
-        if not os.path.exists(filename):
+        if not filename.exists():
             raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), filename)
 
         data: Dict[str, Union[str, int, float]] = {}
@@ -185,7 +189,7 @@ class Metric(ABC):
         return data
 
     @staticmethod
-    def json_write(filename: str, what_to_write: Dict) -> None:
+    def json_write(filename: Path, what_to_write: Dict) -> None:
         """
         Write inside the specified JSON file the mean and stddev.
 
@@ -194,14 +198,14 @@ class Metric(ABC):
             what_to_write (dict): A dictionary containing what to write.
 
         """
-        if os.path.exists(filename):
+        if filename.exists():
             data = Metric.json_read(filename)
             for key in what_to_write:
                 data[key] = str(what_to_write[key])
         else:
             data = what_to_write
-            if not os.path.exists(os.path.dirname(filename)):
-                os.makedirs(os.path.dirname(filename))
+            if not filename.parent.exists():
+                filename.parent.mkdir()
 
         with open(filename, "w+") as fp:
             json.dump(data, fp, indent=4)
diff --git a/src/ashpy/metrics/sliced_wasserstein_metric.py b/src/ashpy/metrics/sliced_wasserstein_metric.py
index 055a124..6fda492 100644
--- a/src/ashpy/metrics/sliced_wasserstein_metric.py
+++ b/src/ashpy/metrics/sliced_wasserstein_metric.py
@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 import operator
-import os
+from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Union
 
 import numpy as np
@@ -28,6 +28,8 @@ from ashpy.metrics.sliced_wasserstein import sliced_wasserstein_distance
 if TYPE_CHECKING:
     from ashpy.contexts import GANContext  # pylint: disable=ungrouped-imports
 
+__ALL__ = ["SingleSWD", "SlicedWassersteinDistance"]
+
 
 class SingleSWD(Metric):
     """SlicedWassersteinDistance for a certain level of the pyramid."""
@@ -35,7 +37,7 @@ class SingleSWD(Metric):
     def __init__(
         self,
         model_selection_operator: Callable = operator.lt,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         level_of_pyramid: int = 0,
         real_or_fake: str = "fake",
     ) -> None:
@@ -94,7 +96,7 @@ class SlicedWassersteinDistance(Metric):
         self,
         name: str = "SWD",
         model_selection_operator: Callable = operator.lt,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         resolution: int = 128,
         resolution_min: int = 16,
         patches_per_image: int = 64,
diff --git a/src/ashpy/metrics/ssim_multiscale.py b/src/ashpy/metrics/ssim_multiscale.py
index 1b68c90..9eee13b 100644
--- a/src/ashpy/metrics/ssim_multiscale.py
+++ b/src/ashpy/metrics/ssim_multiscale.py
@@ -17,8 +17,8 @@ from __future__ import annotations
 
 import math
 import operator
-import os
-from typing import TYPE_CHECKING, Callable, Tuple
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Tuple, Union
 
 import tensorflow as tf
 from ashpy import LogEvalMode
@@ -27,6 +27,8 @@ from ashpy.metrics import Metric
 if TYPE_CHECKING:
     from ashpy.contexts import GANContext  # pylint: disable=ungrouped-imports
 
+__ALL__ = ["SSIM_Multiscale"]
+
 
 class SSIM_Multiscale(Metric):  # pylint: disable=invalid-name
     r"""
@@ -43,11 +45,11 @@ class SSIM_Multiscale(Metric):  # pylint: disable=invalid-name
         self,
         name: str = "SSIM_Multiscale",
         model_selection_operator: Callable = operator.lt,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         max_val: float = 2.0,
         power_factors=None,
         filter_size: int = 11,
-        filter_sigma: int = 1.5,
+        filter_sigma: float = 1.5,
         k1: float = 0.01,
         k2: float = 0.03,
     ) -> None:
@@ -75,7 +77,7 @@ class SSIM_Multiscale(Metric):  # pylint: disable=invalid-name
                 being downsampled by 2.  Defaults to (0.0448, 0.2856, 0.3001, 0.2363,
                 0.1333), which are the values obtained in the original paper.
             filter_size (int): Default value 11 (size of gaussian filter).
-            filter_sigma (int): Default value 1.5 (width of gaussian filter).
+            filter_sigma (float): Default value 1.5 (width of gaussian filter).
             k1 (float): Default value 0.01.
             k2 (float): Default value 0.03 (SSIM is less sensitivity to K2 for lower values, so
                 it would be better if we take the values in range of 0< K2 <0.4).
diff --git a/src/ashpy/restorers/restorer.py b/src/ashpy/restorers/restorer.py
index 9f66a48..518dff2 100644
--- a/src/ashpy/restorers/restorer.py
+++ b/src/ashpy/restorers/restorer.py
@@ -16,7 +16,7 @@
 
 import json
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Optional, Union
 
 import ashpy
 import tensorflow as tf
@@ -32,7 +32,10 @@ class Restorer:
     """
 
     def __init__(
-        self, logdir: str = "log", ckpts_dir: str = "ckpts", expect_partial: bool = True
+        self,
+        logdir: Union[Path, str] = Path().cwd() / "log",
+        ckpts_dir: str = "ckpts",
+        expect_partial: bool = True,
     ) -> None:
         """
         Initialize the Restorer.
@@ -44,15 +47,29 @@ class Restorer:
                 For more information see the docs for :py:func:`tf.train.Checkpoint.restore()`.
 
         """
-        self._ckpts_dir = Path(logdir).joinpath(ckpts_dir)
+        self._ckpts_dir = Path(logdir) / ckpts_dir
         if not self._ckpts_dir.exists():
             raise FileNotFoundError(f"{ckpts_dir} does not exist.")
         self._restored_log_msg = "Restored {} from checkpoint {}."
-        self._human_checkpoint_map = self._read_human_checkpoint_map()
+        try:
+            self._human_checkpoint_map: Optional[
+                Dict[str, str]
+            ] = self._read_human_checkpoint_map()
+        except FileNotFoundError:
+            self._human_checkpoint_map = None
 
     @property
-    def checkpoint_map(self) -> Dict[str, str]:
-        """Get the map of the ids in the checkpoint."""
+    def checkpoint_map(self) -> Optional[Dict[str, str]]:
+        """
+        Get the map of the ids in the checkpoint.
+
+        Map is a Dict where keys are the `ids` in the checkpoint and the values are the
+            string representation of the types.
+
+        Returns:
+            Dict if the map is found, else None.
+
+        """
         return self._human_checkpoint_map
 
     def _restore_checkpoint(self, checkpoint, partial: bool = True):
@@ -129,5 +146,5 @@ class Restorer:
         return callback
 
     def _read_human_checkpoint_map(self) -> Dict[str, str]:
-        with open(self._ckpts_dir.joinpath("checkpoint_map.json")) as fp:
+        with open(self._ckpts_dir / "checkpoint_map.json") as fp:
             return json.load(fp)
diff --git a/src/ashpy/trainers/classifier.py b/src/ashpy/trainers/classifier.py
index 44ce6b1..8f2b2c7 100644
--- a/src/ashpy/trainers/classifier.py
+++ b/src/ashpy/trainers/classifier.py
@@ -14,8 +14,8 @@
 
 """Primitive Trainer Interface."""
 
-import os
-from typing import List, Optional
+from pathlib import Path
+from typing import List, Optional, Union
 
 import ashpy
 import tensorflow as tf
@@ -43,7 +43,7 @@ class ClassifierTrainer(Trainer):
         epochs: int,
         metrics: Optional[List[Metric]] = None,
         callbacks: Optional[List[Callback]] = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         global_step: Optional[tf.Variable] = None,
     ):
         r"""
diff --git a/src/ashpy/trainers/gan.py b/src/ashpy/trainers/gan.py
index 1439328..09e6e2e 100644
--- a/src/ashpy/trainers/gan.py
+++ b/src/ashpy/trainers/gan.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 
 """Collection of GANs trainers."""
-import os
-from typing import List, Optional
+from pathlib import Path
+from typing import List, Optional, Union
 
 import tensorflow as tf
 from ashpy.callbacks import Callback
@@ -140,7 +140,7 @@ class AdversarialTrainer(Trainer):
         epochs: int,
         metrics: Optional[List[Metric]] = None,
         callbacks: Optional[List[Callback]] = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
         global_step: Optional[tf.Variable] = None,
     ):
@@ -389,6 +389,8 @@ class EncoderTrainer(AdversarialTrainer):
     Examples:
         .. testcode::
 
+            from pathlib import Path
+
             import shutil
             import operator
 
@@ -427,9 +429,9 @@ class EncoderTrainer(AdversarialTrainer):
                 [tf.keras.layers.Dense(10), tf.keras.layers.Dense(num_classes)]
             )
 
-            logdir = "testlog/adversarial_encoder"
+            logdir = Path("testlog") / "adversarial_encoder"
 
-            if os.path.exists(logdir):
+            if logdir.exists():
                 shutil.rmtree(logdir)
 
             metrics = [metrics.gan.EncodingAccuracy(classifier)]
@@ -491,7 +493,7 @@ class EncoderTrainer(AdversarialTrainer):
         epochs: int,
         metrics: Optional[List[Metric]] = None,
         callbacks: Optional[List[Callback]] = None,
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
         global_step: Optional[tf.Variable] = None,
     ):
diff --git a/src/ashpy/trainers/trainer.py b/src/ashpy/trainers/trainer.py
index aca69b5..3522247 100644
--- a/src/ashpy/trainers/trainer.py
+++ b/src/ashpy/trainers/trainer.py
@@ -15,8 +15,8 @@
 """Primitive Trainer Interface."""
 
 import json
-import os
 from abc import ABC, abstractmethod
+from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
 import tensorflow as tf
@@ -41,7 +41,7 @@ class Trainer(ABC):
         self,
         epochs: int,
         example_dim: Tuple[int, int],
-        logdir: str = os.path.join(os.getcwd(), "log"),
+        logdir: Union[Path, str] = Path().cwd() / "log",
         log_eval_mode: LogEvalMode = LogEvalMode.TEST,
         global_step: Optional[tf.Variable] = None,
         metrics: Optional[List[Metric]] = None,
@@ -109,22 +109,24 @@ class Trainer(ABC):
             for callback in callbacks:
                 ckpt_dict[callback.name] = callback
 
-        self._logdir = logdir
-        self._ckpts_dir = os.path.join(self._logdir, "ckpts")
+        self._logdir = Path(logdir) if not isinstance(logdir, Path) else logdir
+        self._ckpts_dir = self._logdir / "ckpts"
         self._ckpt_dict = None
         self._checkpoint_map: Dict[str, str] = {}
         self._checkpoints = None
         self._manager = None
         self._update_checkpoint(ckpt_dict)
 
+        # NOTE: as of TensorFlow 2.1.0 pathlib types cannot be converted to tensors automatically.
+        # Explicit conversion to string is required
         self._train_summary_writer = tf.summary.create_file_writer(
-            os.path.join(self._logdir, "train")
+            str(self._logdir / "train")
         )
         self._eval_summary_writer = tf.summary.create_file_writer(
-            os.path.join(self._logdir, "eval")
+            str(self._logdir / "eval")
         )
         self._test_summary_writer = tf.summary.create_file_writer(
-            os.path.join(self._logdir, "test")
+            str(self._logdir / "test")
         )
 
         # Initialize the global batch size to a negative number
@@ -155,8 +157,8 @@ class Trainer(ABC):
         """Generate a human readable map of the id and type mapping in the checkpoint."""
         return {id: str(type(self._ckpt_dict[id])) for id in self._ckpt_dict}
 
-    def _write_checkpoint_map(self):
-        with open(os.path.join(self._ckpts_dir, "checkpoint_map.json"), "w") as fp:
+    def _write_checkpoint_map(self, path):
+        with open(Path(path) / "checkpoint_map.json", "w") as fp:
             json.dump(self._checkpoint_map, fp)
 
     @staticmethod
@@ -251,7 +253,7 @@ class Trainer(ABC):
     def _save(self):
         """Save the current checkpointable object status."""
         checkpoint = self._manager.save()
-        self._write_checkpoint_map()
+        self._write_checkpoint_map(self._ckpts_dir)
         # print is captured from pydoc - deterministic output can be used
         # to run tests.
         print(f"[{self._global_step.numpy()}] Saved checkpoint: {checkpoint}")
@@ -287,7 +289,15 @@ class Trainer(ABC):
     def model_selection(self) -> None:
         """Use the metrics to perform model selection."""
         for metric in self._metrics:
-            metric.model_selection(self._checkpoint, self._global_step)
+            model_selection_ckpt_path: Optional[Path] = metric.model_selection(
+                self._checkpoint, self._global_step
+            )
+            # If model selection has been performed
+            if isinstance(model_selection_ckpt_path, Path):
+                model_selection_ckpt_dir = model_selection_ckpt_path.parent
+                # If we haven't already created the checkpoint_map.json
+                if not (model_selection_ckpt_dir / "checkpoint_map.json").exists():
+                    self._write_checkpoint_map(model_selection_ckpt_dir)
 
     def _measure_performance_if_needed(
         self, example: tf.Tensor, measure_performance_freq: int


Test patch:
diff --git a/conftest.py b/conftest.py
index 443287e..d12e717 100644
--- a/conftest.py
+++ b/conftest.py
@@ -17,6 +17,7 @@
 import operator
 import os
 import shutil
+from pathlib import Path
 
 import pytest
 import tensorflow  # pylint: disable=import-error
@@ -49,19 +50,19 @@ def add_common_namespaces(doctest_namespace):
 @pytest.fixture(scope="function")
 def save_dir():
     """Add the save_dir parameter to tests."""
-    m_save_dir = "testlog/savedir"
+    m_save_dir = Path("testlog/savedir")
 
     # Clean before
-    if os.path.exists(m_save_dir):
+    if m_save_dir.exists():
         shutil.rmtree(m_save_dir)
-        assert not os.path.exists(m_save_dir)
+        assert not m_save_dir.exists()
 
     yield m_save_dir
 
     # teardown
-    if os.path.exists(m_save_dir):
+    if m_save_dir.exists():
         shutil.rmtree(m_save_dir)
-        assert not os.path.exists(m_save_dir)
+        assert not m_save_dir.exists()
 
 
 # ------------------------------------------------------------------------------------
diff --git a/tests/callbacks/test_callbacks.py b/tests/callbacks/test_callbacks.py
index f711b98..b382511 100644
--- a/tests/callbacks/test_callbacks.py
+++ b/tests/callbacks/test_callbacks.py
@@ -14,12 +14,8 @@
 
 """Test LogImageGANCallback."""
 
-from pathlib import Path
-
-import tensorflow as tf
 from ashpy.callbacks import LogImageGANCallback
 from ashpy.callbacks.events import Event
-from tensorflow.python.training.tracking import base
 
 from tests.utils.fake_training_loop import fake_adversarial_training_loop
 
diff --git a/tests/callbacks/test_counter_callback.py b/tests/callbacks/test_counter_callback.py
index 47c1f7e..03f4d41 100644
--- a/tests/callbacks/test_counter_callback.py
+++ b/tests/callbacks/test_counter_callback.py
@@ -15,7 +15,6 @@
 """Test CounterCallback."""
 
 import pytest
-import tensorflow as tf
 from ashpy.callbacks import CounterCallback, Event
 from ashpy.models.gans import ConvDiscriminator, ConvGenerator
 
@@ -66,7 +65,7 @@ def _models():
 def test_counter_callback_multiple_events():
     """Counter Callback should not receive multiple events."""
     with pytest.raises(TypeError):
-        clbk = FakeCounterCallback(
+        FakeCounterCallback(
             event=[Event.ON_EPOCH_END],
             name="TestCounterCallbackMultipleEvents",
             fn=lambda context: print("Bloop"),
diff --git a/tests/callbacks/test_save_callback.py b/tests/callbacks/test_save_callback.py
index c282371..f450609 100644
--- a/tests/callbacks/test_save_callback.py
+++ b/tests/callbacks/test_save_callback.py
@@ -14,6 +14,7 @@
 
 """Test Save Callback."""
 import os
+from pathlib import Path
 from typing import Tuple
 
 import pytest
@@ -33,7 +34,9 @@ INCOMPATIBLE_FORMAT_AND_SUB_FORMAT = [(SaveFormat.MODEL, SaveSubFormat.H5)]
 
 @pytest.mark.parametrize("save_format_and_sub_format", COMPATIBLE_FORMAT_AND_SUB_FORMAT)
 def test_save_callback_compatible(
-    tmpdir, save_format_and_sub_format: Tuple[SaveFormat, SaveSubFormat], save_dir: str,
+    tmpdir,
+    save_format_and_sub_format: Tuple[SaveFormat, SaveSubFormat],
+    save_dir: Path,
 ):
     """Test the integration between callbacks and trainer."""
     save_format, save_sub_format = save_format_and_sub_format
@@ -45,8 +48,7 @@ def test_save_callback_compatible(
 
     for model_dir in save_dirs:
         assert save_format.name() in [
-            x.split(os.path.sep)[-1]
-            for x in os.listdir(os.path.join(save_dir, model_dir))
+            x.split(os.path.sep)[-1] for x in os.listdir(save_dir / model_dir)
         ]
 
 
@@ -54,7 +56,9 @@ def test_save_callback_compatible(
     "save_format_and_sub_format", INCOMPATIBLE_FORMAT_AND_SUB_FORMAT
 )
 def test_save_callback_incompatible(
-    tmpdir, save_format_and_sub_format: Tuple[SaveFormat, SaveSubFormat], save_dir: str,
+    tmpdir,
+    save_format_and_sub_format: Tuple[SaveFormat, SaveSubFormat],
+    save_dir: Path,
 ):
     """Test the integration between callbacks and trainer."""
     save_format, save_sub_format = save_format_and_sub_format
@@ -63,10 +67,10 @@ def test_save_callback_incompatible(
         _test_save_callback_helper(tmpdir, save_format, save_sub_format, save_dir)
 
     # assert no folder has been created
-    assert not os.path.exists(save_dir)
+    assert not save_dir.exists()
 
 
-def _test_save_callback_helper(tmpdir, save_format, save_sub_format, save_dir):
+def _test_save_callback_helper(tmpdir, save_format, save_sub_format, save_dir: Path):
     image_resolution = (28, 28)
     layer_spec_input_res = (7, 7)
     layer_spec_target_res = (7, 7)
@@ -115,23 +119,19 @@ def test_save_callback_type_error(save_dir: str,):
     or save sub-format is passed.
     """
     with pytest.raises(TypeError):
-        callbacks = [
-            SaveCallback(
-                models=[],
-                save_dir=save_dir,
-                verbose=1,
-                save_format="save_format",
-                save_sub_format=SaveSubFormat.TF,
-            )
-        ]
+        SaveCallback(
+            models=[],
+            save_dir=save_dir,
+            verbose=1,
+            save_format="save_format",
+            save_sub_format=SaveSubFormat.TF,
+        )
 
     with pytest.raises(TypeError):
-        callbacks = [
-            SaveCallback(
-                models=[],
-                save_dir=save_dir,
-                verbose=1,
-                save_format=SaveFormat.WEIGHTS,
-                save_sub_format="sub-format",
-            )
-        ]
+        SaveCallback(
+            models=[],
+            save_dir=save_dir,
+            verbose=1,
+            save_format=SaveFormat.WEIGHTS,
+            save_sub_format="sub-format",
+        )
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index c455948..93540a3 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -21,10 +21,8 @@ import json
 import operator
 import pathlib
 import shutil
-from typing import List
 
 import pytest
-import tensorflow as tf
 from ashpy.metrics import ClassifierLoss
 
 from tests.utils.fake_training_loop import fake_classifier_training_loop
@@ -69,9 +67,9 @@ def test_metrics_log(fake_training, tmpdir, cleanup):
     assert not pathlib.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).joinpath("best", metric.sanitized_name)
+        metric_dir = pathlib.Path(tmpdir) / "best" / metric.sanitized_name
         assert metric_dir.exists()
-        json_path = metric_dir.joinpath(f"{metric.sanitized_name}.json")
+        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)
@@ -107,4 +105,4 @@ def test_metrics_names_collision(training_loop, tmpdir):
     """
     metrics = [ClassifierLoss(name="test_loss"), ClassifierLoss(name="test_loss")]
     with pytest.raises(ValueError):
-        training_completed, trainer = training_loop(metrics=metrics, logdir=tmpdir)
+        training_loop(metrics=metrics, logdir=tmpdir)
diff --git a/tests/test_restorers.py b/tests/test_restorers.py
index 2218a1d..c7851a4 100644
--- a/tests/test_restorers.py
+++ b/tests/test_restorers.py
@@ -12,288 +12,296 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Test Restorers."""
+"""
+Test Restorers.
+
+GIVEN a correctly instantiated trainer
+GIVEN some training has been done
+    WHEN calling the Restorer "Restored ... from checkpoint: .../ckpts/ckpt-..."
+        should be logged.
+    WHEN restoring models the first layer of the restored model and the trained one should
+        have the same weights.
+"""
 
-import shutil
 from pathlib import Path
 
 import pytest
 import tensorflow as tf
-from ashpy.restorers import (
-    AdversarialEncoderRestorer,
-    AdversarialRestorer,
-    ClassifierRestorer,
-    Restorer,
-)
-from ashpy.trainers import AdversarialTrainer, ClassifierTrainer, Trainer
+from ashpy.callbacks import CounterCallback
+from ashpy.restorers import AdversarialRestorer, ClassifierRestorer, Restorer
+from ashpy.trainers import AdversarialTrainer, ClassifierTrainer
 
 from tests.utils.fake_datasets import (
     fake_adversarial_dataset,
     fake_autoencoder_datasest,
 )
 from tests.utils.fake_models import basic_dcgan, conv_autoencoder
-from tests.utils.fake_training_loop import fake_classifier_training_loop
 
 DEFAULT_CKPT_DIR = "ckpts"
 
 
-class TestRestorer:
+def _check_first_layer(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)
+
+
+def _test_restore_object(restorer, placeholder, ckpt_id, capsys):
+    """Test that the object is restored correctly."""
+    restorer.restore_object(placeholder, ckpt_id)
+    _check_log(restorer, ckpt_id, capsys)
+
+
+def _check_log(restorer, ckpt_id, capsys):
+    """Test that the object is restored correctly by looking at the logs."""
+    out, _ = capsys.readouterr()
+    # Assert that the log is correct
+    assert restorer._restored_log_msg.format(ckpt_id, restorer._ckpts_dir) in out.split(
+        "\n"
+    )
+
+
+def test_restore_model(fake_training, capsys, tmpdir):
     """
-    Test the Restorer.
-
-    GIVEN a correctly instantiated trainer
-    GIVEN some training has been done
-        WHEN calling the Restorer "Restored ... from checkpoint: .../ckpts/ckpt-..."
-            should be logged.
-        WHEN restoring models the first layer of the restored model and the trained one should
-            have the same weights.
+    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)
 
-    def _check_first_layer(
-        self, 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}")
-            self._check_first_layer(trained, restored, i)
-
-    def _test_restore_object(self, restorer, placeholder, ckpt_id, capsys):
-        """Test that the object is restored correctly."""
-        restorer.restore_object(placeholder, ckpt_id)
-        self._check_log(restorer, ckpt_id, capsys)
-
-    def _check_log(self, restorer, ckpt_id, capsys):
-        """Test that the object is restored correctly by looking at the logs."""
-        out, err = capsys.readouterr()
-        # Assert that the log is correct
-        assert restorer._restored_log_msg.format(
-            ckpt_id, restorer._ckpts_dir
-        ) in out.split("\n")
+    if isinstance(trainer, ClassifierTrainer):
 
-    def test_restore_model(self, fake_training, capsys, tmpdir):
-        """
-        Test that models are correctly restored.
+        placeholder = conv_autoencoder()
 
-        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
+        # Ensure model have been built correctly
+        x, _ = next(iter(fake_autoencoder_datasest()))
+        placeholder(x)
+
+        ckpt_id = trainer.ckpt_id_model
+        _test_restore_object(restorer, placeholder, ckpt_id, capsys)
+        _check_first_layer(trainer._model, placeholder)
+
+    elif isinstance(trainer, AdversarialTrainer):
+
+        placeholder_g, placeholder_d = basic_dcgan(**loop_args)
+
+        # Ensure model have been built correctly
+        (x, _), z = next(iter(fake_adversarial_dataset(**loop_args)))
+        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)
+        _test_restore_object(
+            restorer, placeholder_d, trainer.ckpt_id_discriminator, capsys
         )
-        assert training_completed
-        restorer = Restorer(logdir=tmpdir)
+        _check_first_layer(trainer._discriminator, placeholder_d)
+
 
-        if isinstance(trainer, ClassifierTrainer):
+def test_restore_common_variables(fake_training, capsys, tmpdir):
+    """
+    Test that the convenience methods exposed by :class:`Restorer` work correctly.
 
-            placeholder = conv_autoencoder()
+    The common :class:`tf.Variable`s that can be restored from :class:`Restorer` are:
+        - global step
+        - steps per epoch
 
-            # Ensure model have been built correctly
-            x, y = next(iter(fake_autoencoder_datasest()))
-            placeholder(x)
+    """
+    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)
+
+    # Restore variables and check their values using the convenience method
+    assert tf.equal(trainer._global_step, restorer.get_global_step())
+    assert tf.equal(trainer._steps_per_epoch, restorer.get_steps_per_epoch())
+
+    out, _ = capsys.readouterr()
+
+    # Check the log
+    for id_to_check in [
+        trainer.ckpt_id_global_step,
+        trainer.ckpt_id_steps_per_epoch,
+    ]:
+        # Assert that the log is correct
+        assert restorer._restored_log_msg.format(
+            id_to_check, restorer._ckpts_dir
+        ) in out.split("\n")
 
-            ckpt_id = trainer.ckpt_id_model
-            self._test_restore_object(restorer, placeholder, ckpt_id, capsys)
-            self._check_first_layer(trainer._model, placeholder)
 
-        elif isinstance(trainer, AdversarialTrainer):
+def test_restore_callbacks(fake_training, 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)
+
+    if isinstance(trainer, AdversarialTrainer):
+        placeholder_callbacks = loop_args.get("callbacks")
+        for i, placeholder_callback in enumerate(placeholder_callbacks):
+            # Restore the callbacks
+            restorer.restore_callback(placeholder_callback, placeholder_callback.name)
+            # Check the log
+            _check_log(restorer, placeholder_callback.name, capsys)
+            # Check that the trained values and the restored are equal
+            if isinstance(placeholder_callback, CounterCallback):
+                assert tf.equal(
+                    placeholder_callbacks[i]._event_counter,
+                    trainer._callbacks[i]._event_counter,
+                )
 
-            placeholder_g, placeholder_d = basic_dcgan(**loop_args)
 
-            # Ensure model have been built correctly
-            (x, y), z = next(iter(fake_adversarial_dataset(**loop_args)))
-            fake = placeholder_g(z)
-            assert tf.reduce_all(tf.equal(fake.shape, x.shape))
-            placeholder_d(x)
+def test_read_checkpoint_map(fake_training, 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)
+    assert restorer.checkpoint_map == trainer._generate_checkpoint_map()
+
+    # Test that Restorer.checkpoint_map without the checkpoint_map.json correctly returns None
+    # Remove checkpoint_map.json
+    ckpt_map: Path = restorer._ckpts_dir / "checkpoint_map.json"
+    ckpt_map.unlink()
+    assert not ckpt_map.exists()
+    assert not Restorer(tmpdir).checkpoint_map
+
+
+# ###################################################
+# Test Convenience Methods
+
+
+def _test_convenience_model_restorer(
+    restorer: AdversarialRestorer,
+    convenience_method,
+    placeholder_model,
+    trained_model,
+    ckpt_id,
+    capsys,
+):
+    convenience_method(placeholder_model)
+    _check_log(restorer, ckpt_id, capsys)
+    _check_first_layer(trained_model, placeholder_model)
+
+
+def _test_convenience_optimizer_restorer(
+    restorer, convenience_method, placeholder_optimizer, ckpt_id, capsys
+):
+    """
+    Test that the various optimizers are correctly restored using convenience classes.
 
-            self._test_restore_object(
-                restorer, placeholder_g, trainer.ckpt_id_generator, capsys
-            )
-            self._check_first_layer(trainer._generator, placeholder_g)
-            self._test_restore_object(
-                restorer, placeholder_d, trainer.ckpt_id_discriminator, capsys
-            )
-            self._check_first_layer(trainer._discriminator, placeholder_d)
+    TODO: Add a more thorough check like :meth:`_check_first_layer()`
+    """
+    convenience_method(placeholder_optimizer)
+    _check_log(restorer, ckpt_id, capsys)
 
-    def test_restore_common_variables(self, fake_training, capsys, tmpdir):
-        """
-        Test that the convenience methods exposed by :class:`Restorer` work correctly.
 
-        The common :class:`tf.Variable`s that can be restored from :class:`Restorer` are:
-            - global step
-            - steps per epoch
+def test_convenience_restorer(fake_training, capsys, tmpdir):
+    """
+    Test that models and optimizers are correctly restored using the convenience classes.
 
-        """
-        training_loop, loop_args, metrics = fake_training
-        training_completed, trainer = training_loop(
-            logdir=tmpdir, metrics=metrics, **loop_args
+    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
+
+    if isinstance(trainer, ClassifierTrainer):
+        restorer: ClassifierRestorer = ClassifierRestorer(logdir=logdir)
+
+        placeholder_model = conv_autoencoder()
+        placeholder_opt = tf.keras.optimizers.Adam()
+
+        # Ensure model have been built correctly
+        x, _ = next(iter(fake_autoencoder_datasest()))
+        placeholder_model(x)
+
+        _test_convenience_model_restorer(
+            restorer,
+            restorer.restore_model,
+            placeholder_model,
+            trainer._model,
+            trainer.ckpt_id_model,
+            capsys,
         )
-        assert training_completed
-        restorer = Restorer(logdir=tmpdir)
-
-        # Restore variables and check their values using the convenience method
-        assert tf.equal(trainer._global_step, restorer.get_global_step())
-        assert tf.equal(trainer._steps_per_epoch, restorer.get_steps_per_epoch())
-
-        out, err = capsys.readouterr()
-
-        # Check the log
-        for id_to_check in [
-            trainer.ckpt_id_global_step,
-            trainer.ckpt_id_steps_per_epoch,
-        ]:
-            # Assert that the log is correct
-            assert restorer._restored_log_msg.format(
-                id_to_check, restorer._ckpts_dir
-            ) in out.split("\n")
-
-    def test_restore_callbacks(self, fake_training, 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
+        _test_convenience_optimizer_restorer(
+            restorer,
+            restorer.restore_optimizer,
+            placeholder_opt,
+            trainer.ckpt_id_optimizer,
+            capsys,
         )
-        assert training_completed
-        restorer = Restorer(logdir=tmpdir)
-
-        if isinstance(trainer, AdversarialTrainer):
-            placeholder_callbacks = loop_args.get("callbacks")
-            for placeholder_callback in placeholder_callbacks:
-                restorer.restore_callback(
-                    placeholder_callback, placeholder_callback._name
-                )
 
-    def test_read_checkpoint_map(self, fake_training, capsys, 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
+    elif isinstance(trainer, AdversarialTrainer):
+        restorer: AdversarialRestorer = AdversarialRestorer(logdir=logdir)
+
+        placeholder_g, placeholder_d = basic_dcgan(**loop_args)
+        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)))
+        fake = placeholder_g(z)
+        assert tf.reduce_all(tf.equal(fake.shape, x.shape))
+        placeholder_d(x)
+
+        _test_convenience_model_restorer(
+            restorer,
+            restorer.restore_generator,
+            placeholder_g,
+            trainer._generator,
+            trainer.ckpt_id_generator,
+            capsys,
+        )
+        _test_convenience_optimizer_restorer(
+            restorer,
+            restorer.restore_generator_optimizer,
+            placeholder_optimizer_g,
+            trainer.ckpt_id_optimizer_generator,
+            capsys,
+        )
+        _test_convenience_model_restorer(
+            restorer,
+            restorer.restore_discriminator,
+            placeholder_d,
+            trainer._discriminator,
+            trainer.ckpt_id_discriminator,
+            capsys,
         )
-        assert training_completed
-        restorer = Restorer(logdir=tmpdir)
-        assert restorer.checkpoint_map == trainer._generate_checkpoint_map()
-
-    # ###################################################
-    # Test Convenience Methods
-
-    def _test_convenience_model_restorer(
-        self,
-        restorer: AdversarialRestorer,
-        convenience_method,
-        placeholder_model,
-        trained_model,
-        ckpt_id,
-        capsys,
-    ):
-        convenience_method(placeholder_model)
-        self._check_log(restorer, ckpt_id, capsys)
-        self._check_first_layer(trained_model, placeholder_model)
-
-    def _test_convenience_optimizer_restorer(
-        self, restorer, convenience_method, placeholder_optimizer, ckpt_id, capsys
-    ):
-        """
-        Test that the various optimizers are correctly restored using convenience classes.
-
-        TODO: Add a more thorough check like :meth:`_check_first_layer()`
-        """
-        convenience_method(placeholder_optimizer)
-        self._check_log(restorer, ckpt_id, capsys)
-
-    def test_convenience_restorer(self, fake_training, 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
+        _test_convenience_optimizer_restorer(
+            restorer,
+            restorer.restore_discriminator_optimizer,
+            placeholder_optimizer_d,
+            trainer.ckpt_id_optimizer_discriminator,
+            capsys,
         )
-        assert training_completed
-
-        if isinstance(trainer, ClassifierTrainer):
-            restorer: ClassifierRestorer = ClassifierRestorer(logdir=logdir)
-
-            placeholder_model = conv_autoencoder()
-            placeholder_opt = tf.keras.optimizers.Adam()
-
-            # Ensure model have been built correctly
-            x, y = next(iter(fake_autoencoder_datasest()))
-            placeholder_model(x)
-
-            ckpt_id = trainer.ckpt_id_model
-            self._test_convenience_model_restorer(
-                restorer,
-                restorer.restore_model,
-                placeholder_model,
-                trainer._model,
-                trainer.ckpt_id_model,
-                capsys,
-            )
-            self._test_convenience_optimizer_restorer(
-                restorer,
-                restorer.restore_optimizer,
-                placeholder_opt,
-                trainer.ckpt_id_optimizer,
-                capsys,
-            )
-
-        elif isinstance(trainer, AdversarialTrainer):
-            restorer: AdversarialRestorer = AdversarialRestorer(logdir=logdir)
-
-            placeholder_g, placeholder_d = basic_dcgan(**loop_args)
-            placeholder_optimizer_g, placeholder_optimizer_d = (
-                tf.keras.optimizers.Adam(),
-                tf.keras.optimizers.Adam(),
-            )
-
-            # Ensure model have been built correctly
-            (x, y), z = next(iter(fake_adversarial_dataset(**loop_args)))
-            fake = placeholder_g(z)
-            assert tf.reduce_all(tf.equal(fake.shape, x.shape))
-            placeholder_d(x)
-
-            self._test_convenience_model_restorer(
-                restorer,
-                restorer.restore_generator,
-                placeholder_g,
-                trainer._generator,
-                trainer.ckpt_id_generator,
-                capsys,
-            )
-            self._test_convenience_optimizer_restorer(
-                restorer,
-                restorer.restore_generator_optimizer,
-                placeholder_optimizer_g,
-                trainer.ckpt_id_optimizer_generator,
-                capsys,
-            )
-            self._test_convenience_model_restorer(
-                restorer,
-                restorer.restore_discriminator,
-                placeholder_d,
-                trainer._discriminator,
-                trainer.ckpt_id_discriminator,
-                capsys,
-            )
-            self._test_convenience_optimizer_restorer(
-                restorer,
-                restorer.restore_discriminator_optimizer,
-                placeholder_optimizer_g,
-                trainer.ckpt_id_optimizer_discriminator,
-                capsys,
-            )
 
 
 def test_failings(tmpdir):
@@ -305,7 +313,7 @@ def test_failings(tmpdir):
     # Test Restorer fails on empty checkpoint dir
     with pytest.raises(FileNotFoundError):
         restorer = Restorer(tmpdir)
-        status = restorer._restore_checkpoint(tf.train.Checkpoint())
+        restorer._restore_checkpoint(tf.train.Checkpoint())
 
     # Test failed placeholders validation
     with pytest.raises(TypeError):
diff --git a/tests/test_trainers.py b/tests/test_trainers.py
index fd264ef..d0dd6cb 100644
--- a/tests/test_trainers.py
+++ b/tests/test_trainers.py
@@ -15,6 +15,9 @@
 """Tests for :mod:`ashpy.trainers`."""
 
 from pathlib import Path
+from typing import List
+
+import ashpy
 
 
 def test_generate_human_ckpt_dict(fake_training, tmpdir):
@@ -23,9 +26,14 @@ def test_generate_human_ckpt_dict(fake_training, tmpdir):
 
     TODO: improve the test.
     """
+    logdir = Path(tmpdir)
     training_loop, loop_args, metrics = fake_training
     training_completed, trainer = training_loop(
-        logdir=tmpdir, metrics=metrics, **loop_args
+        logdir=logdir, metrics=metrics, **loop_args
     )
+    assert training_completed
     assert trainer._checkpoint_map
-    assert Path(trainer._ckpts_dir).joinpath("checkpoint_map.json").exists()
+    assert (Path(trainer._ckpts_dir) / "checkpoint_map.json").exists()
+    metrics: List[ashpy.metrics.Metric] = trainer._metrics
+    for metric in metrics:
+        assert metric.best_folder / "ckpts" / "checkpoint_map.json"
diff --git a/tests/utils/fake_training_loop.py b/tests/utils/fake_training_loop.py
index 8abee78..74ce155 100644
--- a/tests/utils/fake_training_loop.py
+++ b/tests/utils/fake_training_loop.py
@@ -14,6 +14,8 @@
 
 """Fake training loop to simplify training in tests."""
 import operator
+from pathlib import Path
+from typing import Union
 
 import ashpy
 import tensorflow as tf
@@ -32,7 +34,7 @@ __ALL__ = ["fake_classifier_training_loop", "fake_adversarial_training_loop"]
 
 def fake_classifier_training_loop(
     # Trainer
-    logdir: str = "testlog",
+    logdir: Union[Path, str] = "testlog",
     optimizer=tf.optimizers.Adam(1e-4),
     metrics=[ashpy.metrics.ClassifierLoss(model_selection_operator=operator.lt)],
     epochs=2,
@@ -88,7 +90,7 @@ def fake_classifier_training_loop(
 
 
 def fake_adversarial_training_loop(
-    logdir: str = "testlog",
+    logdir: Union[Path, str] = "testlog",
     kernel_size=(5, 5),
     metrics=None,
     callbacks=None,