Repository: pytorch/ignite

Issue:
Backwards compability for torch.save in ModelCheckpoint
## 🚀 Feature

<!-- A clear and concise description of the feature proposal -->
Pytorch made some changes to torch.save (https://pytorch.org/docs/stable/generated/torch.save.html)

> The 1.6 release of PyTorch switched torch.save to use a new zipfile-based file format. torch.load still retains the ability to load files in the old format. If for any reason you want torch.save to use the old format, pass the kwarg _use_new_zipfile_serialization=False.

It would be useful to have control over this flag when instantiating ModelCheckpoint.

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
<!-- A clear and concise description of what you want to happen. -->
<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
Possible ways that I can think of:
- Have some *args, **kwargs passed on to the save function
- Be able to pass the save function itself to ModelCheckpoints
<!-- Add any other context or screenshots about the feature request here. -->


Discussion:
kilsenp: With the current approach it seems to me that it's a bit hard to test if the keyword is passed correctly to the save_fn.

I added therefore two tests
- One to check if this one specific keyword _use_new_zipfile_serialization keyword is accepted (introduced in torch 1.4.0)
- One to check if the keywords are passed correctly by passing an unknown keyword and expecting it to fail.


Patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
index 67e33002..155877ec 100644
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -526,14 +526,16 @@ class DiskSaver(BaseSaveHandler):
         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``.
+        **kwargs: Accepted keyword arguments for `torch.save` or `xm.save`.
     """
 
     def __init__(
-        self, dirname: str, atomic: bool = True, create_dir: bool = True, require_empty: bool = True,
+        self, dirname: str, atomic: bool = True, create_dir: bool = True, require_empty: bool = True, **kwargs
     ):
         self.dirname = os.path.expanduser(dirname)
         self._atomic = atomic
         self._check_and_setup(dirname, create_dir, require_empty)
+        self.kwargs = kwargs
 
     @staticmethod
     @idist.one_rank_only()
@@ -575,7 +577,7 @@ class DiskSaver(BaseSaveHandler):
 
     def _save_func(self, checkpoint: Mapping, path: str, func: Callable, rank: int = 0):
         if not self._atomic:
-            func(checkpoint, path)
+            func(checkpoint, path, **self.kwargs)
         else:
             tmp_file = None
             tmp_name = None
@@ -585,7 +587,7 @@ class DiskSaver(BaseSaveHandler):
                 tmp_file = tmp.file
                 tmp_name = tmp.name
             try:
-                func(checkpoint, tmp_file)
+                func(checkpoint, tmp_file, **self.kwargs)
             except BaseException:
                 if tmp is not None:
                     tmp.close()
@@ -653,6 +655,7 @@ class ModelCheckpoint(Checkpoint):
         archived (bool, optional): Deprecated argument as models saved by `torch.save` are already compressed.
         include_self (bool): Whether to include the `state_dict` of this object in the checkpoint. If `True`, then
             there must not be another object in ``to_save`` with key ``checkpointer``.
+        **kwargs: Accepted keyword arguments for `torch.save` or `xm.save` in `DiskSaver`.
 
     Examples:
         >>> import os
@@ -685,6 +688,7 @@ class ModelCheckpoint(Checkpoint):
         global_step_transform: Optional[Callable] = None,
         archived: bool = False,
         include_self: bool = False,
+        **kwargs
     ):
 
         if not save_as_state_dict:
@@ -704,7 +708,7 @@ class ModelCheckpoint(Checkpoint):
                 # No choice
                 raise ValueError(msg)
 
-        disk_saver = DiskSaver(dirname, atomic=atomic, create_dir=create_dir, require_empty=require_empty,)
+        disk_saver = DiskSaver(dirname, atomic=atomic, create_dir=create_dir, require_empty=require_empty, **kwargs)
 
         super(ModelCheckpoint, self).__init__(
             to_save=None,


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
index 0369cc27..6478ff49 100644
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -6,6 +6,7 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 import torch.nn as nn
+from pkg_resources import parse_version
 
 import ignite.distributed as idist
 from ignite.engine import Engine, Events, State
@@ -592,6 +593,31 @@ def test_disk_saver_atomic(dirname):
     _test_existance(atomic=True, _to_save=to_save_non_serializable, expected=False)
 
 
+@pytest.mark.skipif(
+    parse_version(torch.__version__) < parse_version("1.4.0"), reason="Zipfile serialization was introduced in 1.4.0"
+)
+def test_disk_saver_zipfile_serialization_keyword(dirname):
+    model = DummyModel()
+    to_save = {"model": model}
+
+    saver = DiskSaver(dirname, create_dir=False, _use_new_zipfile_serialization=False)
+    fname = "test.pt"
+    saver(to_save, fname)
+    fp = os.path.join(saver.dirname, fname)
+    assert os.path.exists(fp)
+    saver.remove(fname)
+
+
+def test_disk_saver_unknown_keyword(dirname):
+    model = DummyModel()
+    to_save = {"model": model}
+
+    saver = DiskSaver(dirname, create_dir=False, unknown_keyword="")
+    fname = "test.pt"
+    with pytest.raises(TypeError, match=r"got an unexpected keyword argument 'unknown_keyword'"):
+        saver(to_save, fname)
+
+
 def test_last_k(dirname):
 
     h = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=2)