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 thanks for the suggestion ! 

Maybe, passing args/kwargs would make sense. We also have `Checkpoint` class where user can specify how to save checkpoints. For example, for local storage we can use `DiskSaver` which calls `torch.save`. Forwarding `args/kwargs` to `DiskSaver` can be possible...
Ah okay, I see.
So when going the route of passing kwargs, I would envision it like this:
```
class DiskSaver:
def __init__(..., **additional_kwargs):
    self.additional_kwargs = additional_kwargs

def _save_func(...):
    ...
    func(..., **self.additional_kwargs)
```
and for ModelCheckpoint, the constructor is changed as well.

I can create a pull request for this if this is a good way to go forward.


> I can create a pull request for this if this is a good way to go forward.

@kilsenp yes, I think it can be a good way to to do that and, yes, please, send a PR. The only thing we have to check and deal somehow is [`xm.save`](https://github.com/pytorch/xla/blob/0c34871b91159e51603119c42d9fe58fb7d1da88/torch_xla/core/xla_model.py#L788) from pytorch/xla. Either we do not pass kwargs there or we have to open an Issue/PR at pytorch/xla and propagate things there too... We can start with the 1st option and add 2nd once pytorch/xla are agreed on that.

Resolving patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
--- 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 @@ def _save_xla(self, checkpoint: Mapping, path: str):
 
     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 @@ def _save_func(self, checkpoint: Mapping, path: str, func: Callable, rank: int =
                 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 @@ def __init__(
         global_step_transform: Optional[Callable] = None,
         archived: bool = False,
         include_self: bool = False,
+        **kwargs
     ):
 
         if not save_as_state_dict:
@@ -704,7 +708,7 @@ def __init__(
                 # 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
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -6,6 +6,7 @@
 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_existance(atomic, _to_save, expected):
     _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)