Repository: pytorch/ignite

Issue:
Custom configuration how/where to save the best model with ModelCheckpoint
In case when user would like to integrate `ModelCheckpoint` with a package for experiment tracking, e.g. mlflow, polyaxon, etc. 
In such case logging, model weights etc can be stored on a cloud storage, e.g.
```python
exp_tracking.log_artifact(filepath)
```
Be default `ModelCheckpoint` is saving the model to the provided path `dirname`.
Idea is to provide a flexibility to execute a custom code when model is saved and be eable to store everywhere we would like.

What do you think ?

cc @elanmart 


Discussion:
I have been using a hackish solution for this problem. A lambda function at the end of the function _internal_save() that receives the path to the saved file.
 This work fine to use with the mlflow.log_artifact function that expects a file. However I few that this approach is somewhat hackish. Maybe a better approach would be to use handles to attach "output targets".
This sounds nice. I also did some hacking and just want to share the code, in case it can be useful for anyone of you. In my case, I needed a custom save method. I didn't want to use `torch.save()`, because my own model class is still under development and I want to have compatibility between all its versions. My save method simply saves hyperparameters and weights from which the class is recreated when it is loaded. 

I inherited my *ModelSaver* class from `ignite.handlers.ModelCheckpoint` and  overloaded the `_internal_save` method.  There are two more little features: Save the model on exception and when the training is completed. Sorry for the incomplete documentation, but at the moment, I have only little time.

Long story short... here is my code:
```python
from ignite.engine import Events
import ignite
import os


class ModelSaver(ignite.handlers.ModelCheckpoint):
    """"
    Extends class`ignite.handlers.ModelCheckpoint with option to provide a custom save method,
    saving the final model after training ends and saving a model if an exception is raised during training
    """

    def __init__(self, *args, save_method=None, save_on_exception=True, save_on_completed=True, **kwargs):
        if not isinstance(save_on_exception, bool):
            raise TypeError(
                "Argument save_on_exception must be of type bool, got {] instead.".format(type(save_on_exception)))
        if not isinstance(save_on_completed, bool):
            raise TypeError(
                "Argument save_on_completed must be of type bool, got {] instead.".format(type(save_on_completed)))
        if save_method is not None and not callable(save_method):
            raise TypeError(
                "Argument save_method must be callable accepting two arguments: Model object to save and path")

        self._save_method = save_method
        self._save_on_completed = save_on_completed
        self._save_on_exception = save_on_exception

        super(ModelSaver, self).__init__(*args, **kwargs)

    def _internal_save(self, obj, path):
        if self._save_method is not None:
            self._save_method(obj, path)
        else:
            super(ModelSaver, self)._internal_save(obj, path)

    def _on_exception(self, engine, exception, to_save):
        for name, obj in to_save.items():
            fname = '{}_{}_{}{}.pth'.format(self._fname_prefix, name, self._iteration, "_on_exception")
            path = os.path.join(self._dirname, fname)
            if os.path.exists(path):
                os.remove(path)
            self._save(obj=obj, path=path)

    def _on_completed(self, engine, to_save):
        for name, obj in to_save.items():
            fname = '{}_{}_{}{}.pth'.format(self._fname_prefix, name, self._iteration, "_on_completed")
            path = os.path.join(self._dirname, fname)
            if os.path.exists(path):
                os.remove(path)
            self._save(obj=obj, path=path)

    def attach(self, engine, model_dict):
        """
                Attaches the model saver to an engine object

                Args:
                    engine (Engine): engine object
                    model_dict (dict): A dict mapping names to objects, e.g. {'mymodel': model}
        """
        engine.add_event_handler(Events.EPOCH_COMPLETED, self, model_dict)
        engine.add_event_handler(Events.COMPLETED, self._on_completed, model_dict)
        engine.add_event_handler(Events.EXCEPTION_RAISED, self._on_exception, model_dict)

```
@Bibonaut thanks for sharing the code ! Looks nice! We can think to put it into `contrib` module.
@Bibonaut would you like to send a PR with this code and some basic tests ? It would be awesome !
@vfdev-5  Sorry for my late response, still very busy :| I'ld really like to contribute if you can wait until mid of May where I have to meet a deadline. Then I'ld like to share and discuss some more features with you I made to make my workflow and investigations a lot easier. 
@Bibonaut no problems, we can wait until your deadline :)

Resolving patch:
diff --git a/ignite/contrib/handlers/base_logger.py b/ignite/contrib/handlers/base_logger.py
--- a/ignite/contrib/handlers/base_logger.py
+++ b/ignite/contrib/handlers/base_logger.py
@@ -105,7 +105,7 @@ def __init__(self, tag, metric_names=None, output_transform=None, another_engine
             if not isinstance(another_engine, Engine):
                 raise TypeError("Argument another_engine should be of type Engine, "
                                 "but given {}".format(type(another_engine)))
-            warnings.warn("Use of another_engine is deprecated and will be removed in 0.3.0. "
+            warnings.warn("Use of another_engine is deprecated and will be removed in 0.4.0. "
                           "Please use global_step_transform instead.", DeprecationWarning)
             global_step_transform = global_step_from_engine(another_engine)
 
diff --git a/ignite/handlers/__init__.py b/ignite/handlers/__init__.py
--- a/ignite/handlers/__init__.py
+++ b/ignite/handlers/__init__.py
@@ -1,4 +1,4 @@
-from ignite.handlers.checkpoint import ModelCheckpoint
+from ignite.handlers.checkpoint import ModelCheckpoint, Checkpoint, DiskSaver
 from ignite.handlers.timing import Timer
 from ignite.handlers.early_stopping import EarlyStopping
 from ignite.handlers.terminate_on_nan import TerminateOnNan
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -1,11 +1,229 @@
 import os
 import tempfile
+from collections import namedtuple
+import warnings
+
+import sys
+
+IS_PYTHON2 = sys.version_info[0] < 3
+
+if IS_PYTHON2:
+    import collections
+else:
+    import collections.abc as collections
 
 import torch
 
 
-class ModelCheckpoint(object):
-    """ ModelCheckpoint handler can be used to periodically save objects to disk.
+class Checkpoint(object):
+    """Checkpoint handler can be used to periodically save and load objects which have attribute
+    `state_dict`/`load_state_dict`. This class can use specific save handlers to store on the disk or a cloud
+    storage, etc.
+
+    Args:
+        to_save (dict): Dictionary with the objects to save. Objects should have implemented `state_dict` and `
+            load_state_dict` methods.
+        save_handler (callable): Method to use to save engine and other provided objects. Function receives a
+            checkpoint as a dictionary to save. In case if user needs to save engine's checkpoint on a disk,
+            `save_handler` can be defined with :class:`~ignite.handlers.DiskSaver`.
+        filename_prefix (str, optional): Prefix for the filename to which objects will be saved. See Note for details.
+        score_function (callable, optional): If not None, it should be a function taking a single argument,
+            :class:`~ignite.engine.Engine` object, and returning a score (`float`). Objects with highest scores will be
+            retained. Exactly one of (`save_interval`, `score_function`) arguments must be provided.
+        score_name (str, optional): If `score_function` not None, it is possible to store its absolute value using
+            `score_name`. See Notes for more details.
+        n_saved (int, optional): Number of objects that should be kept on disk. Older files will be removed.
+
+    Note:
+        This class stores a single file as a dictionary of provided objects to save.
+        The filename has the following structure: `{filename_prefix}_{name}_{suffix}.pth` where
+
+        - `filename_prefix` is the argument passed to the constructor,
+        - `name` is the key in `to_save` if a single object is to store, otherwise `name` is "checkpoint".
+        - `suffix` can have 3 possible values:
+
+        1) Default, when no `score_function`/`score_name` defined: Suffix is defined by attached engine's current
+        iteration. The filename will be `{filename_prefix}_{name}_{engine.state.iteration}.pth`.
+
+        2) With `score_function`, but without `score_name`: Suffix is defined by provided score. The filename will be
+        `{filename_prefix}_{name}_{score}.pth`.
+
+        3) With `score_function` and `score_name`: In this case, user can store its absolute value using `score_name`
+        in the filename. Suffix is defined by engine's current epoch, score name and its value. The filename will
+        be `{filename_prefix}_{name}_{engine.state.epoch}_{score_name}={abs(score)}.pth`.
+        For example, `score_name="val_loss"` and `score_function` that returns `-loss` (as objects with
+        highest scores will be retained), then saved models filenames will be `model_0_val_loss=0.1234.pth`.
+
+    Examples:
+        Saved checkpoint then can be used to resume a training:
+
+        .. code-block:: python
+
+            from ignite.engine import Engine, Events
+            from ignite.handlers import Checkpoint, DiskSaver
+            from torch import nn
+            trainer = Engine(lambda batch: None)
+            model = ...
+            optimizer = ...
+            lr_scheduler = ...
+            to_save = {'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}
+            handler = Checkpoint(to_save, DiskSaver('/tmp/models', create_dir=True), 'myprefix', n_saved=2)
+            trainer.add_event_handler(Events.EPOCH_COMPLETED, handler)
+            trainer.run([0], max_epochs=6)
+            os.listdir('/tmp/models')
+
+
+    """
+
+    Item = namedtuple("Item", ["priority", "filename"])
+
+    def __init__(self, to_save, save_handler, filename_prefix="",
+                 score_function=None, score_name=None,
+                 n_saved=1):
+
+        if not isinstance(to_save, collections.Mapping):
+            raise TypeError("Argument `to_save` should be a dictionary, but given {}".format(type(to_save)))
+
+        if len(to_save) < 1:
+            raise ValueError("No objects to checkpoint.")
+
+        if not callable(save_handler):
+            raise TypeError("Argument `save_handler` should be callable")
+
+        if score_function is None and score_name is not None:
+            raise ValueError("If `score_name` is provided, then `score_function` "
+                             "should be also provided.")
+
+        self._check_objects(to_save)
+        self._fname_prefix = filename_prefix + "_" if len(filename_prefix) > 0 else filename_prefix
+        self.save_handler = save_handler
+        self.to_save = to_save
+        self._score_function = score_function
+        self._score_name = score_name
+        self._n_saved = n_saved
+        self._saved = []
+
+    @property
+    def last_checkpoint(self):
+        if len(self._saved) < 1:
+            return None
+        return self._saved[0].filename
+
+    def __call__(self, engine):
+
+        if self._score_function is not None:
+            priority = self._score_function(engine)
+        else:
+            priority = engine.state.iteration
+
+        if len(self._saved) < self._n_saved or \
+                self._saved[0].priority < priority:
+
+            suffix = "{}".format(priority)
+            if self._score_name is not None:
+                suffix = "{}_{}={:.7}".format(engine.state.epoch, self._score_name, abs(priority))
+
+            checkpoint = self._setup_checkpoint()
+
+            name = "checkpoint"
+            if len(checkpoint) == 1:
+                for k in checkpoint:
+                    name = k
+            filename = '{}{}_{}.pth'.format(self._fname_prefix, name, suffix)
+
+            self.save_handler(checkpoint, filename)
+
+            self._saved.append(Checkpoint.Item(priority, filename))
+            self._saved.sort(key=lambda item: item[0])
+
+        if len(self._saved) > self._n_saved:
+            item = self._saved.pop(0)
+            self.save_handler.remove(item.filename)
+
+    def _setup_checkpoint(self):
+        checkpoint = {}
+        for k, obj in self.to_save.items():
+            checkpoint[k] = obj.state_dict()
+        return checkpoint
+
+    @staticmethod
+    def _check_objects(to_save_or_load):
+        for k, obj in to_save_or_load.items():
+            if not (hasattr(obj, "state_dict") and hasattr(obj, "load_state_dict")):
+                raise TypeError("Object {} should have `state_dict` and `load_state_dict` methods".format(type(obj)))
+
+    @staticmethod
+    def load_objects(to_load, checkpoint):
+        """Method to apply `load_state_dict` on the objects from `to_load` using states from `checkpoint`.
+
+        Args:
+            to_load (Mapping):
+            checkpoint (Mapping):
+        """
+        Checkpoint._check_objects(to_load)
+        if not isinstance(checkpoint, collections.Mapping):
+            raise TypeError("Argument checkpoint should be a dictionary, but given {}".format(type(checkpoint)))
+        for k, obj in to_load.items():
+            if k not in checkpoint:
+                raise ValueError("Object labeled by '{}' from `to_load` is not found in the checkpoint".format(k))
+            obj.load_state_dict(checkpoint[k])
+
+
+class DiskSaver(object):
+    """Handler that saves input checkpoint on a disk.
+
+    Args:
+        dirname (str): Directory path where the checkpoint will be saved
+        atomic (bool, optional): if True, checkpoint is serialized to a temporary file, and then
+            moved to final destination, so that files are guaranteed to not be damaged
+            (for example if exception occures during saving).
+        create_dir (bool, optional): if True, will create directory 'dirname' if it doesnt exist.
+        require_empty (bool, optional): If True, will raise exception if there are any files in the directory 'dirname'.
+
+    """
+
+    def __init__(self, dirname, atomic=True, create_dir=True, require_empty=True):
+        self.dirname = os.path.expanduser(dirname)
+        self._atomic = atomic
+        if create_dir:
+            if not os.path.exists(dirname):
+                os.makedirs(dirname)
+        # Ensure that dirname exists
+        if not os.path.exists(dirname):
+            raise ValueError("Directory path '{}' is not found".format(dirname))
+
+        if require_empty:
+            matched = [fname for fname in os.listdir(dirname)]
+            if len(matched) > 0:
+                raise ValueError("Files are already present "
+                                 "in the directory {}. If you want to use this "
+                                 "directory anyway, pass `require_empty=False`."
+                                 "".format(dirname))
+
+    def __call__(self, checkpoint, filename):
+        path = os.path.join(self.dirname, filename)
+
+        if not self._atomic:
+            torch.save(checkpoint, path)
+        else:
+            tmp = tempfile.NamedTemporaryFile(delete=False, dir=self.dirname)
+            try:
+                torch.save(checkpoint, tmp.file)
+            except BaseException:
+                tmp.close()
+                os.remove(tmp.name)
+                raise
+            else:
+                tmp.close()
+                os.rename(tmp.name, path)
+
+    def remove(self, filename):
+        path = os.path.join(self.dirname, filename)
+        os.remove(path)
+
+
+class ModelCheckpoint(Checkpoint):
+    """ModelCheckpoint handler can be used to periodically save objects to disk.
 
     This handler expects two arguments:
 
@@ -14,15 +232,28 @@ class ModelCheckpoint(object):
 
     See Notes and Examples for further details.
 
+    .. warning::
+        Behaviour of this class has been changed since v0.3.0.
+
+        Argument `save_as_state_dict` is deprecated and should not be used. It is considered as True.
+
+        Argument `save_interval` is deprecated and should not be used. Please, use events filtering instead, e.g.
+        :attr:`~ignite.engine.Events.ITERATION_STARTED(every=1000)`
+
+        There is no more internal counter that has been used to indicate the number of save actions. User could
+        see its value `step_number` in the filename, e.g. `{filename_prefix}_{name}_{step_number}.pth`. Actually,
+        `step_number` is replaced by current engine's epoch if `score_function` is specified and current iteration
+        otherwise.
+
+        A single `pth` file is created instead of multiple files.
+
+
     Args:
         dirname (str):
             Directory path where objects will be saved.
         filename_prefix (str):
-            Prefix for the filenames to which objects will be saved. See Notes
+            Prefix for the filenames to which objects will be saved. See Notes of :class:`~ignite.handlers.Checkpoint`
             for more details.
-        save_interval (int, optional):
-            if not None, objects will be saved to disk every `save_interval` calls to the handler.
-            Exactly one of (`save_interval`, `score_function`) arguments must be provided.
         score_function (callable, optional):
             if not None, it should be a function taking a single argument,
             an :class:`~ignite.engine.Engine` object,
@@ -42,25 +273,6 @@ class ModelCheckpoint(object):
             in the directory 'dirname'.
         create_dir (bool, optional):
             If True, will create directory 'dirname' if it doesnt exist.
-        save_as_state_dict (bool, optional):
-            If True, will save only the `state_dict` of the objects specified, otherwise the whole object will be saved.
-
-    Note:
-          This handler expects two arguments: an :class:`~ignite.engine.Engine` object and a `dict`
-          mapping names to objects that should be saved.
-
-          These names are used to specify filenames for saved objects.
-          Each filename has the following structure:
-          `{filename_prefix}_{name}_{step_number}.pth`.
-          Here, `filename_prefix` is the argument passed to the constructor,
-          `name` is the key in the aforementioned `dict`, and `step_number`
-          is incremented by `1` with every call to the handler.
-
-          If `score_function` is provided, user can store its absolute value using `score_name` in the filename.
-          Each filename can have the following structure:
-          `{filename_prefix}_{name}_{step_number}_{score_name}={abs(score_function_result)}.pth`.
-          For example, `score_name="val_loss"` and `score_function` that returns `-loss` (as objects with highest scores
-          will be retained), then saved models filenames will be `model_resnet_10_val_loss=0.1234.pth`.
 
     Examples:
         >>> import os
@@ -68,114 +280,59 @@ class ModelCheckpoint(object):
         >>> from ignite.handlers import ModelCheckpoint
         >>> from torch import nn
         >>> trainer = Engine(lambda batch: None)
-        >>> handler = ModelCheckpoint('/tmp/models', 'myprefix', save_interval=2, n_saved=2, create_dir=True)
+        >>> handler = ModelCheckpoint('/tmp/models', 'myprefix', n_saved=2, create_dir=True)
         >>> model = nn.Linear(3, 3)
-        >>> trainer.add_event_handler(Events.EPOCH_COMPLETED, handler, {'mymodel': model})
+        >>> trainer.add_event_handler(Events.EPOCH_COMPLETED(every=2), handler, {'mymodel': model})
         >>> trainer.run([0], max_epochs=6)
         >>> os.listdir('/tmp/models')
         ['myprefix_mymodel_4.pth', 'myprefix_mymodel_6.pth']
     """
 
     def __init__(self, dirname, filename_prefix,
-                 save_interval=None, score_function=None, score_name=None,
+                 save_interval=None,
+                 score_function=None, score_name=None,
                  n_saved=1,
                  atomic=True, require_empty=True,
                  create_dir=True,
                  save_as_state_dict=True):
 
-        self._dirname = os.path.expanduser(dirname)
-        self._fname_prefix = filename_prefix
-        self._n_saved = n_saved
-        self._save_interval = save_interval
-        self._score_function = score_function
-        self._score_name = score_name
-        self._atomic = atomic
-        self._saved = []  # list of tuples (priority, saved_objects)
-        self._iteration = 0
-        self._save_as_state_dict = save_as_state_dict
+        if not save_as_state_dict:
+            raise ValueError("Argument save_as_state_dict is deprecated and should be True")
+        if save_interval is not None:
+            msg = "Argument save_interval is deprecated and should be None. " \
+                  "Please, use events filtering instead, e.g. Events.ITERATION_STARTED(every=1000)"
+            if save_interval == 1:
+                # Do not break for old version who used `save_interval=1`
+                warnings.warn(msg)
+            else:
+                # No choice
+                raise ValueError(msg)
 
-        if not (save_interval is None) ^ (score_function is None):
-            raise ValueError("Exactly one of `save_interval`, or `score_function` "
-                             "arguments must be provided.")
+        disk_saver = DiskSaver(dirname, atomic=atomic, create_dir=create_dir, require_empty=require_empty)
 
         if score_function is None and score_name is not None:
             raise ValueError("If `score_name` is provided, then `score_function` "
                              "should be also provided.")
 
-        if create_dir:
-            if not os.path.exists(dirname):
-                os.makedirs(dirname)
-
-        # Ensure that dirname exists
-        if not os.path.exists(dirname):
-            raise ValueError("Directory path '{}' is not found.".format(dirname))
-
-        if require_empty:
-            matched = [fname
-                       for fname in os.listdir(dirname)
-                       if fname.startswith(self._fname_prefix)]
-
-            if len(matched) > 0:
-                raise ValueError("Files prefixed with {} are already present "
-                                 "in the directory {}. If you want to use this "
-                                 "directory anyway, pass `require_empty=False`."
-                                 "".format(filename_prefix, dirname))
-
-    def _save(self, obj, path):
-        if not self._atomic:
-            self._internal_save(obj, path)
-        else:
-            tmp = tempfile.NamedTemporaryFile(delete=False, dir=self._dirname)
-            try:
-                self._internal_save(obj, tmp.file)
-            except BaseException:
-                tmp.close()
-                os.remove(tmp.name)
-                raise
-            else:
-                tmp.close()
-                os.rename(tmp.name, path)
+        self._fname_prefix = filename_prefix + "_" if len(filename_prefix) > 0 else filename_prefix
+        self.save_handler = disk_saver
+        self.to_save = None
+        self._score_function = score_function
+        self._score_name = score_name
+        self._n_saved = n_saved
+        self._saved = []
 
-    def _internal_save(self, obj, path):
-        if not self._save_as_state_dict:
-            torch.save(obj, path)
-        else:
-            if not hasattr(obj, "state_dict") or not callable(obj.state_dict):
-                raise ValueError("Object should have `state_dict` method.")
-            torch.save(obj.state_dict(), path)
+    @property
+    def last_checkpoint(self):
+        if len(self._saved) < 1:
+            return None
+        return os.path.join(self.save_handler.dirname, self._saved[0].filename)
 
     def __call__(self, engine, to_save):
+
         if len(to_save) == 0:
             raise RuntimeError("No objects to checkpoint found.")
 
-        self._iteration += 1
-
-        if self._score_function is not None:
-            priority = self._score_function(engine)
-
-        else:
-            priority = self._iteration
-            if (self._iteration % self._save_interval) != 0:
-                return
-
-        if (len(self._saved) < self._n_saved) or (self._saved[0][0] < priority):
-            saved_objs = []
-
-            suffix = ""
-            if self._score_name is not None:
-                suffix = "_{}={:.7}".format(self._score_name, abs(priority))
-
-            for name, obj in to_save.items():
-                fname = '{}_{}_{}{}.pth'.format(self._fname_prefix, name, self._iteration, suffix)
-                path = os.path.join(self._dirname, fname)
-
-                self._save(obj=obj, path=path)
-                saved_objs.append(path)
-
-            self._saved.append((priority, saved_objs))
-            self._saved.sort(key=lambda item: item[0])
-
-        if len(self._saved) > self._n_saved:
-            _, paths = self._saved.pop(0)
-            for p in paths:
-                os.remove(p)
+        self._check_objects(to_save)
+        self.to_save = to_save
+        super(ModelCheckpoint, self).__call__(engine)


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -1,13 +1,16 @@
 import os
 import tempfile
+import shutil
+import warnings
 
-import pytest
 import torch
 import torch.nn as nn
-import shutil
 
-from ignite.engine import Engine, Events
-from ignite.handlers import ModelCheckpoint
+from ignite.engine import Engine, Events, State
+from ignite.handlers import ModelCheckpoint, Checkpoint, DiskSaver
+
+import pytest
+from mock import MagicMock
 
 _PREFIX = 'PREFIX'
 
@@ -38,32 +41,35 @@ def test_args_validation(dirname):
     with open(os.path.join(nonempty, '{}_name_0.pth'.format(_PREFIX)), 'w'):
         pass
 
-    # save_interval & score_func
-    with pytest.raises(ValueError):
-        h = ModelCheckpoint(existing, _PREFIX,
-                            create_dir=False)
+    with pytest.raises(ValueError, match=r"Files are already present in the directory"):
+        ModelCheckpoint(nonempty, _PREFIX)
 
-    with pytest.raises(ValueError):
-        h = ModelCheckpoint(nonempty, _PREFIX, save_interval=42)
+    with pytest.raises(ValueError, match=r"Argument save_interval is deprecated and should be None"):
+        ModelCheckpoint(existing, _PREFIX, save_interval=42)
 
-    with pytest.raises(ValueError):
-        h = ModelCheckpoint(nonempty, _PREFIX,
-                            score_function="score",
-                            save_interval=42)
+    with pytest.raises(ValueError, match=r"Directory path '\S+' is not found"):
+        ModelCheckpoint(os.path.join(dirname, 'non_existing_dir'), _PREFIX, create_dir=False)
 
-    with pytest.raises(ValueError):
-        h = ModelCheckpoint(os.path.join(dirname, 'non_existing_dir'), _PREFIX,
-                            create_dir=False,
-                            save_interval=42)
+    with pytest.raises(ValueError, match=r"Argument save_as_state_dict is deprecated and should be True"):
+        ModelCheckpoint(existing, _PREFIX, create_dir=False, save_as_state_dict=False)
 
 
 def test_simple_recovery(dirname):
-    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, save_interval=1,
-                        save_as_state_dict=False)
-    h(None, {'obj': 42})
+    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False)
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=1)
+
+    model = DummyModel()
+    to_save = {'model': model}
+    h(engine, to_save)
 
-    fname = os.path.join(dirname, '{}_{}_{}.pth'.format(_PREFIX, 'obj', 1))
-    assert torch.load(fname) == 42
+    fname = h.last_checkpoint
+    assert isinstance(fname, str)
+    assert os.path.join(dirname, _PREFIX) in fname
+    assert os.path.exists(fname)
+    loaded_objects = torch.load(fname)
+    assert "model" in loaded_objects
+    assert loaded_objects['model'] == model.state_dict()
 
 
 def test_simple_recovery_from_existing_non_empty(dirname):
@@ -71,72 +77,90 @@ def test_simple_recovery_from_existing_non_empty(dirname):
     with open(previous_fname, 'w') as f:
         f.write("test")
 
-    h = ModelCheckpoint(dirname, _PREFIX, create_dir=True, require_empty=False,
-                        save_interval=1, save_as_state_dict=False)
-    h(None, {'obj': 42})
+    h = ModelCheckpoint(dirname, _PREFIX, create_dir=True, require_empty=False)
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=1)
 
-    fname = os.path.join(dirname, '{}_{}_{}.pth'.format(_PREFIX, 'obj', 1))
-    assert torch.load(fname) == 42
+    model = DummyModel()
+    to_save = {'model': model}
+    h(engine, to_save)
+
+    fname = h.last_checkpoint
+    assert isinstance(fname, str)
+    assert os.path.join(dirname, '{}_{}_{}.pth'.format(_PREFIX, 'model', 1)) == fname
+    assert os.path.exists(fname)
     assert os.path.exists(previous_fname)
+    loaded_objects = torch.load(fname)
+    assert "model" in loaded_objects
+    assert loaded_objects['model'] == model.state_dict()
 
 
 def test_atomic(dirname):
-    serializable = 42
-    non_serializable = (42, lambda _: 42)
 
-    def _test_existance(atomic, name, obj, expected):
-        h = ModelCheckpoint(dirname, _PREFIX,
-                            atomic=atomic,
-                            create_dir=False,
-                            require_empty=False,
-                            save_interval=1,
-                            save_as_state_dict=False)
+    model = DummyModel()
+    to_save_serializable = {'model': model}
+    to_save_non_serializable = {'model': lambda x: x}
+
+    def _test_existance(atomic, _to_save, expected):
 
+        saver = DiskSaver(dirname, atomic=atomic, create_dir=False, require_empty=False)
+        fname = "test.pth"
         try:
-            h(None, {name: obj})
+            with warnings.catch_warnings():
+                # Ignore torch/serialization.py:292: UserWarning: Couldn't retrieve source code for container of type
+                # DummyModel. It won't be checked for correctness upon loading.
+                warnings.simplefilter("ignore", category=UserWarning)
+                saver(_to_save, fname)
         except Exception:
             pass
+        fp = os.path.join(saver.dirname, fname)
+        assert os.path.exists(fp) == expected
+        if expected:
+            saver.remove(fname)
 
-        fname = os.path.join(dirname, '{}_{}_{}.pth'.format(_PREFIX, name, 1))
-        assert os.path.exists(fname) == expected
-
-    _test_existance(atomic=False, name='nonatomic_OK', obj=serializable, expected=True)
-    _test_existance(atomic=False, name='nonatomic_FAIL', obj=non_serializable, expected=True)
+    _test_existance(atomic=False, _to_save=to_save_serializable, expected=True)
+    _test_existance(atomic=False, _to_save=to_save_non_serializable, expected=True)
 
-    _test_existance(atomic=True, name='atomic_OK', obj=serializable, expected=True)
-    _test_existance(atomic=True, name='atomic_FAIL', obj=non_serializable, expected=False)
+    _test_existance(atomic=True, _to_save=to_save_serializable, expected=True)
+    _test_existance(atomic=True, _to_save=to_save_non_serializable, expected=False)
 
 
 def test_last_k(dirname):
-    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=2,
-                        save_interval=2, save_as_state_dict=False)
-    to_save = {'name': 42}
 
-    for _ in range(8):
-        h(None, to_save)
+    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=2)
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=0)
+
+    model = DummyModel()
+    to_save = {'model': model}
+    h(engine, to_save)
 
-    expected = ['{}_{}_{}.pth'.format(_PREFIX, 'name', i)
-                for i in [6, 8]]
+    for i in range(1, 9):
+        engine.state.iteration = i
+        h(engine, to_save)
 
-    assert sorted(os.listdir(dirname)) == expected
+    expected = ['{}_{}_{}.pth'.format(_PREFIX, 'model', i) for i in [7, 8]]
+
+    assert sorted(os.listdir(dirname)) == expected, "{} vs {}".format(sorted(os.listdir(dirname)), expected)
 
 
 def test_best_k(dirname):
-    scores = iter([1.0, -2., 3.0, -4.0])
+    scores = iter([1.2, -2., 3.1, -4.0])
 
-    def score_function(engine):
+    def score_function(_):
         return next(scores)
 
-    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False,
-                        n_saved=2, score_function=score_function,
-                        save_as_state_dict=False)
+    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=2, score_function=score_function)
 
-    to_save = {'name': 42}
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=0)
+
+    model = DummyModel()
+    to_save = {'model': model}
     for _ in range(4):
-        h(None, to_save)
+        h(engine, to_save)
 
-    expected = ['{}_{}_{}.pth'.format(_PREFIX, 'name', i)
-                for i in [1, 3]]
+    expected = ['{}_{}_{}.pth'.format(_PREFIX, 'model', i) for i in [1.2, 3.1]]
 
     assert sorted(os.listdir(dirname)) == expected
 
@@ -149,92 +173,52 @@ def score_function(engine):
         return next(scores_iter)
 
     h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=2,
-                        score_function=score_function, score_name="val_loss",
-                        save_as_state_dict=False)
+                        score_function=score_function, score_name="val_loss")
 
-    to_save = {'name': 42}
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=0)
+
+    model = DummyModel()
+    to_save = {'model': model}
     for _ in range(4):
-        h(None, to_save)
+        engine.state.epoch += 1
+        h(engine, to_save)
 
-    expected = ['{}_{}_{}_val_loss={:.7}.pth'.format(_PREFIX, 'name', i, scores[i - 1])
-                for i in [1, 3]]
+    expected = ['{}_{}_{}_val_loss={:.7}.pth'.format(_PREFIX, 'model', e, scores[e - 1]) for e in [1, 3]]
 
     assert sorted(os.listdir(dirname)) == expected
 
 
 def test_with_engine(dirname):
 
-    def update_fn(engine, batch):
+    def update_fn(_1, _2):
         pass
 
     name = 'model'
     engine = Engine(update_fn)
-    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False,
-                              n_saved=2, save_interval=1,
-                              save_as_state_dict=False)
+    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=2)
 
-    engine.add_event_handler(Events.EPOCH_COMPLETED, handler, {name: 42})
+    model = DummyModel()
+    to_save = {'model': model}
+    engine.add_event_handler(Events.EPOCH_COMPLETED, handler, to_save)
     engine.run([0], max_epochs=4)
 
-    expected = ['{}_{}_{}.pth'.format(_PREFIX, name, i)
-                for i in [3, 4]]
+    expected = ['{}_{}_{}.pth'.format(_PREFIX, name, i) for i in [3, 4]]
 
     assert sorted(os.listdir(dirname)) == expected
 
 
-def test_no_state_dict(dirname):
-    model = DummyModel()
+def test_with_state_dict(dirname):
 
-    def update_fn(engine, batch):
+    def update_fn(_1, _2):
         pass
 
-    name = 'model'
     engine = Engine(update_fn)
-    handler = ModelCheckpoint(
-        dirname,
-        _PREFIX,
-        create_dir=False,
-        n_saved=1,
-        save_interval=1,
-        save_as_state_dict=False)
-
-    engine.add_event_handler(Events.EPOCH_COMPLETED, handler, {name: model})
-    engine.run([0], max_epochs=4)
+    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
 
-    saved_model = os.path.join(dirname, os.listdir(dirname)[0])
-    load_model = torch.load(saved_model)
-
-    assert isinstance(load_model, DummyModel)
-    assert not isinstance(load_model, dict)
-
-    model_state_dict = model.state_dict()
-    loaded_model_state_dict = load_model.state_dict()
-    for key in model_state_dict.keys():
-        assert key in loaded_model_state_dict
-
-        model_value = model_state_dict[key]
-        loaded_model_value = loaded_model_state_dict[key]
-
-        assert model_value.numpy() == loaded_model_value.numpy()
-
-
-def test_with_state_dict(dirname):
     model = DummyModel()
-
-    def update_fn(engine, batch):
-        pass
-
-    name = 'model'
-    engine = Engine(update_fn)
-    handler = ModelCheckpoint(
-        dirname,
-        _PREFIX,
-        create_dir=False,
-        n_saved=1,
-        save_interval=1,
-        save_as_state_dict=True)
-
-    engine.add_event_handler(Events.EPOCH_COMPLETED, handler, {name: model})
+    to_save = {'model': model}
+    engine.add_event_handler(Events.EPOCH_COMPLETED, handler, to_save)
     engine.run([0], max_epochs=4)
 
     saved_model = os.path.join(dirname, os.listdir(dirname)[0])
@@ -246,42 +230,27 @@ def update_fn(engine, batch):
     model_state_dict = model.state_dict()
     loaded_model_state_dict = load_model
     for key in model_state_dict.keys():
-        assert key in loaded_model_state_dict
+        assert key in loaded_model_state_dict['model']
 
         model_value = model_state_dict[key]
-        loaded_model_value = loaded_model_state_dict[key]
+        loaded_model_value = loaded_model_state_dict['model'][key]
 
         assert model_value.numpy() == loaded_model_value.numpy()
 
 
 def test_valid_state_dict_save(dirname):
     model = DummyModel()
-    h = ModelCheckpoint(
-        dirname,
-        _PREFIX,
-        create_dir=False,
-        n_saved=1,
-        save_interval=1,
-        save_as_state_dict=True)
+    h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
+
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=0)
 
     to_save = {'name': 42}
-    with pytest.raises(ValueError):
-        h(None, to_save)
-    to_save = {'name': "string"}
-    with pytest.raises(ValueError):
-        h(None, to_save)
-    to_save = {'name': []}
-    with pytest.raises(ValueError):
-        h(None, to_save)
-    to_save = {'name': {}}
-    with pytest.raises(ValueError):
-        h(None, to_save)
-    to_save = {'name': ()}
-    with pytest.raises(ValueError):
-        h(None, to_save)
+    with pytest.raises(TypeError, match=r"should have `state_dict` and `load_state_dict` methods"):
+        h(engine, to_save)
     to_save = {'name': model}
     try:
-        h(None, to_save)
+        h(engine, to_save)
     except ValueError:
         pytest.fail("Unexpected ValueError")
 
@@ -301,13 +270,7 @@ def update_fn(engine, batch):
         lr_scheduler.step()
 
     engine = Engine(update_fn)
-    handler = ModelCheckpoint(
-        dirname,
-        _PREFIX,
-        create_dir=False,
-        n_saved=1,
-        save_interval=1,
-        save_as_state_dict=True)
+    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
 
     engine.add_event_handler(Events.EPOCH_COMPLETED,
                              handler,
@@ -319,14 +282,15 @@ def update_fn(engine, batch):
     engine.run([0], max_epochs=4)
 
     saved_objects = sorted(os.listdir(dirname))
-    # saved objects are ['PREFIX_lr_scheduler_4.pth', 'PREFIX_model_4.pth', 'PREFIX_optimizer_4.pth']
-    saved_lr_scheduler = os.path.join(dirname, saved_objects[0])
-    saved_model = os.path.join(dirname, saved_objects[1])
-    saved_optimizer = os.path.join(dirname, saved_objects[2])
+    # saved object is ['PREFIX_checkpoint_4.pth', ]
+    saved_checkpoint = os.path.join(dirname, saved_objects[0])
 
-    loaded_model_state_dict = torch.load(saved_model)
-    loaded_optimizer_state_dict = torch.load(saved_optimizer)
-    loaded_lr_scheduler_state_dict = torch.load(saved_lr_scheduler)
+    loaded_obj = torch.load(saved_checkpoint)
+    for f in ["model", "optimizer", "lr_scheduler"]:
+        assert f in loaded_obj
+    loaded_model_state_dict = loaded_obj['model']
+    loaded_optimizer_state_dict = loaded_obj['optimizer']
+    loaded_lr_scheduler_state_dict = loaded_obj['lr_scheduler']
 
     assert isinstance(loaded_model_state_dict, dict)
     assert isinstance(loaded_optimizer_state_dict, dict)
@@ -352,3 +316,160 @@ def update_fn(engine, batch):
         lr_scheduler_value = lr_scheduler_state_dict[key]
         loaded_lr_scheduler_value = loaded_lr_scheduler_state_dict[key]
         assert lr_scheduler_value == loaded_lr_scheduler_value
+
+
+def test_checkpoint_wrong_input():
+
+    with pytest.raises(TypeError, match=r"Argument `to_save` should be a dictionary"):
+        Checkpoint(12, lambda x: x, "prefix", )
+
+    with pytest.raises(TypeError, match=r"Argument `to_save` should be a dictionary"):
+        Checkpoint([12, ], lambda x: x, "prefix")
+
+    with pytest.raises(ValueError, match=r"No objects to checkpoint."):
+        Checkpoint({}, lambda x: x, "prefix")
+
+    model = DummyModel()
+    to_save = {'model': model}
+
+    with pytest.raises(TypeError, match=r"Argument `save_handler` should be callable"):
+        Checkpoint(to_save, 12, "prefix")
+
+    with pytest.raises(ValueError,
+                       match=r"If `score_name` is provided, then `score_function` should be also provided."):
+        Checkpoint(to_save, lambda x: x, score_name="acc")
+
+
+def test_checkpoint__setup_checkpoint():
+    save_handler = MagicMock()
+
+    to_save = {'model1': DummyModel(), 'model2': DummyModel()}
+
+    checkpointer = Checkpoint(to_save, save_handler=save_handler)
+    chkpt = checkpointer._setup_checkpoint()
+    assert isinstance(chkpt, dict)
+    for k in ['model1', 'model2']:
+        assert k in chkpt
+        assert chkpt[k] == to_save[k].state_dict()
+
+
+def test_checkpoint_default():
+    save_handler = MagicMock()
+    save_handler.remove = MagicMock()
+
+    model = DummyModel()
+    to_save = {'model': model}
+
+    checkpointer = Checkpoint(to_save, save_handler=save_handler)
+
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=0)
+    checkpointer(engine)
+
+    assert save_handler.call_count == 1
+    save_handler.assert_called_with(checkpointer._setup_checkpoint(), "model_0.pth")
+
+    engine.state.iteration += 1
+    checkpointer(engine)
+    assert save_handler.call_count == 2
+    save_handler.assert_called_with(checkpointer._setup_checkpoint(), "model_1.pth")
+    assert save_handler.remove.call_count == 1
+    save_handler.remove.assert_called_with("model_0.pth")
+
+
+def test_checkpoint_score_function():
+
+    model = DummyModel()
+    to_save = {'model': model}
+
+    def _test(scores, score_name=None):
+        save_handler = MagicMock()
+        save_handler.remove = MagicMock()
+
+        def score_function(_):
+            return next(scores)
+
+        checkpointer = Checkpoint(to_save, save_handler=save_handler,
+                                  score_function=score_function, score_name=score_name)
+
+        engine = Engine(lambda e, b: None)
+        engine.state = State(epoch=0, iteration=0)
+        checkpointer(engine)
+
+        assert save_handler.call_count == 1
+        fname = "model_{}1.2.pth"
+        if score_name is not None:
+            fname = fname.format("0_" + score_name + "=")
+        else:
+            fname = fname.format("")
+        save_handler.assert_called_with(checkpointer._setup_checkpoint(), fname)
+
+        engine.state.epoch += 1
+        checkpointer(engine)
+        assert save_handler.call_count == 1
+
+        engine.state.epoch += 1
+        checkpointer(engine)
+        assert save_handler.call_count == 2
+        fname2 = "model_{}3.1.pth"
+        if score_name is not None:
+            fname2 = fname2.format("2_" + score_name + "=")
+        else:
+            fname2 = fname2.format("")
+        save_handler.assert_called_with(checkpointer._setup_checkpoint(), fname2)
+        assert save_handler.remove.call_count == 1
+        save_handler.remove.assert_called_with(fname)
+
+    _test(iter([1.2, -2., 3.1, -4.0]))
+    _test(iter([1.2, float('nan'), 3.1, -4.0]))
+    _test(iter([1.2, -2., 3.1, -4.0]), score_name="acc")
+    _test(iter([1.2, float('nan'), 3.1, -4.0]), score_name="acc")
+
+
+def test_checkpoint_last_checkpoint():
+    save_handler = MagicMock()
+    save_handler.__call__ = MagicMock()
+    model = DummyModel()
+    to_save = {'model': model}
+
+    checkpointer = Checkpoint(to_save, save_handler=save_handler)
+    assert checkpointer.last_checkpoint is None
+
+    engine = Engine(lambda e, b: None)
+    engine.state = State(epoch=0, iteration=0)
+
+    checkpointer(engine)
+    assert checkpointer.last_checkpoint == "model_0.pth"
+
+
+def test_checkpoint_load_objects():
+
+    with pytest.raises(TypeError, match=r"Argument checkpoint should be a dictionary"):
+        Checkpoint.load_objects({}, [])
+
+    model = DummyModel()
+    to_load = {'model': model}
+
+    with pytest.raises(ValueError, match=r"from `to_load` is not found in the checkpoint"):
+        Checkpoint.load_objects(to_load, {})
+
+    model = DummyModel()
+    to_load = {'model': model}
+    model2 = DummyModel()
+
+    chkpt = {'model': model2.state_dict()}
+    Checkpoint.load_objects(to_load, chkpt)
+    assert model.state_dict() == model2.state_dict()
+
+
+def test_disksaver_wrong_input(dirname):
+
+    with pytest.raises(ValueError, match=r"Directory path '\S+' is not found"):
+        DiskSaver("/tmp/non-existing-folder", create_dir=False)
+
+    previous_fname = os.path.join(dirname, '{}_{}_{}.pth'.format(_PREFIX, 'obj', 1))
+    with open(previous_fname, 'w') as f:
+        f.write("test")
+
+    with pytest.raises(ValueError, match=r"Files are already present in the directory"):
+        DiskSaver(dirname, require_empty=True)