address comments

opacus/grad_sample/grad_sample_module.py

@@ -433,26 +433,19 @@ def state_dict(self, *args, **kwargs) -> Dict:
             f"_module.{key}": value
             for key, value in self._module.state_dict(*args, **kwargs).items()
         }
-        ret_state_dict["batch_first"] = self.batch_first
-        ret_state_dict["loss_reduction"] = self.loss_reduction
         return ret_state_dict
 
     def load_state_dict(self, state_dict: Dict, **kwargs):
         """
         Load the state_dict into the wrapped module
         """
         state_dict = state_dict.copy()
-        self.batch_first = state_dict.pop("batch_first", self.batch_first)
-        self.loss_reduction = state_dict.pop("loss_reduction", self.loss_reduction)
         # remove "_module." prefix before loading into wrapped module
         for key in list(state_dict.keys()):
             if key.startswith("_module."):
                 prefix_stripped_key = key[len("_module.") :]
                 state_dict[prefix_stripped_key] = state_dict.pop(key)
-        self._module.load_state_dict(state_dict, **kwargs)
-        # remove and add hooks with the newly loaded loss_reduction and batch_first
-        self.remove_hooks()
-        self.add_hooks(loss_reduction=self.loss_reduction, batch_first=self.batch_first)
+        return self._module.load_state_dict(state_dict, **kwargs)
 
     @classmethod
     def is_supported(cls, module: nn.Module) -> bool:

opacus/privacy_engine.py

@@ -503,6 +503,7 @@ def save_checkpoint(
         module: GradSampleModule,
         optimizer: Optional[DPOptimizer] = None,
         noise_scheduler: Optional[_NoiseScheduler] = None,
+        checkpoint_dict: Optional[Dict[str, Any]] = None,
         module_state_dict_kwargs: Optional[Dict[str, Any]] = None,
         torch_save_kwargs: Optional[Dict[str, Any]] = None,
     ):
@@ -516,17 +517,17 @@ def save_checkpoint(
             torch_save_kwargs: dict of kwargs to pass to ``torch.save()``
 
         """
-        dict_to_save = {}
-        dict_to_save["module_state_dict"] = module.state_dict(
+        checkpoint_dict = checkpoint_dict or {}
+        checkpoint_dict["module_state_dict"] = module.state_dict(
             **(module_state_dict_kwargs or {})
         )
-        dict_to_save["privacy_accountant_state_dict"] = self.accountant.state_dict()
+        checkpoint_dict["privacy_accountant_state_dict"] = self.accountant.state_dict()
         if optimizer is not None:
-            dict_to_save["optimizer_state_dict"] = optimizer.state_dict()
+            checkpoint_dict["optimizer_state_dict"] = optimizer.state_dict()
         if noise_scheduler is not None:
-            dict_to_save["noise_scheduler_state_dict"] = noise_scheduler.state_dict()
+            checkpoint_dict["noise_scheduler_state_dict"] = noise_scheduler.state_dict()
 
-        torch.save(dict_to_save, path, **(torch_save_kwargs or {}))
+        torch.save(checkpoint_dict, path, **(torch_save_kwargs or {}))
 
     def load_checkpoint(
         self,
@@ -537,7 +538,7 @@ def load_checkpoint(
         noise_scheduler: Optional[_NoiseScheduler] = None,
         module_load_dict_kwargs: Optional[Dict[str, Any]] = None,
         torch_load_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> Dict:
         checkpoint = torch.load(path, **(torch_load_kwargs or {}))
         module.load_state_dict(
             checkpoint["module_state_dict"], **(module_load_dict_kwargs or {})
@@ -547,7 +548,15 @@ def load_checkpoint(
         optimizer_state_dict = checkpoint.pop("optimizer_state_dict", {})
         if optimizer is not None and len(optimizer_state_dict) > 0:
             optimizer.load_state_dict(optimizer_state_dict)
+        elif (optimizer is not None) ^ (len(optimizer_state_dict) > 0):
+            # warn if only one of them is available
+            warnings.warn(
+                f"optimizer_state_dict has {len(optimizer_state_dict)} items"
+                f" but optimizer is {'' if optimizer else 'not'} provided."
+            )
 
         noise_scheduler_state_dict = checkpoint.pop("noise_scheduler_state_dict", {})
         if noise_scheduler is not None and len(noise_scheduler_state_dict) > 0:
             noise_scheduler.load_state_dict(noise_scheduler_state_dict)
+
+        return checkpoint

opacus/tests/grad_sample_module_test.py

@@ -232,32 +232,17 @@ def test_submodule_access(self):
     def test_state_dict(self):
         gs_state_dict = self.grad_sample_module.state_dict()
         og_state_dict = self.original_model.state_dict()
-        # check gs specific attributes
-        self.assertEqual(gs_state_dict["batch_first"], True)
-        self.assertEqual(gs_state_dict["loss_reduction"], "mean")
         # check wrapped module state dict
         for key in og_state_dict.keys():
             self.assertTrue(f"_module.{key}" in gs_state_dict)
             assert_allclose(og_state_dict[key], gs_state_dict[f"_module.{key}"])
 
     def test_load_state_dict(self):
         gs_state_dict = self.grad_sample_module.state_dict()
-        gs_state_dict["loss_reduction"] = "sum"
-        _ = gs_state_dict.pop("batch_first")
-
         new_gs = GradSampleModule(
             SampleConvNet(), batch_first=False, loss_reduction="mean"
         )
-
-        new_gs_hook_before_load = new_gs.autograd_grad_sample_hooks[0]
         new_gs.load_state_dict(gs_state_dict)
-        new_gs_hook_after_load = new_gs.autograd_grad_sample_hooks[0]
-
-        self.assertEqual(new_gs.loss_reduction, "sum")  # value should have changed
-        self.assertEqual(new_gs.batch_first, False)  # old value to be retained
-        self.assertTrue(
-            new_gs_hook_before_load != new_gs_hook_after_load
-        )  # hook is reset
         # wrapped module is the same
         for key in self.original_model.state_dict().keys():
             self.assertTrue(key in new_gs._module.state_dict())

opacus/tests/privacy_engine_test.py

@@ -491,9 +491,7 @@ def test_checkpoints(self, noise_scheduler: Optional[_NoiseScheduler]):
         )
 
         # check that two sets of components are different
-        self.assertFalse(
-            are_state_dict_equal(m1._module.state_dict(), m2._module.state_dict())
-        )
+        self.assertFalse(are_state_dict_equal(m1.state_dict(), m2.state_dict()))
         if noise_scheduler:
             self.assertNotEqual(s1.state_dict(), s2.state_dict())
         self.assertNotEqual(opt1.noise_multiplier, opt2.noise_multiplier)
@@ -504,20 +502,27 @@ def test_checkpoints(self, noise_scheduler: Optional[_NoiseScheduler]):
             s1.step()
 
         # load into set 2
+        checkpoint_to_save = {"foo": "bar"}
         with io.BytesIO() as bytesio:
             pe1.save_checkpoint(
-                path=bytesio, module=m1, optimizer=opt1, noise_scheduler=s1
+                path=bytesio,
+                module=m1,
+                optimizer=opt1,
+                noise_scheduler=s1,
+                checkpoint_dict=checkpoint_to_save,
             )
             bytesio.seek(0)
-            pe2.load_checkpoint(
+            loaded_checkpoint = pe2.load_checkpoint(
                 path=bytesio, module=m2, optimizer=opt2, noise_scheduler=s2
             )
 
-        # check the two sets of components are now the same
-        self.assertEqual(pe1.accountant.state_dict(), pe2.accountant.state_dict())
+        # check if loaded checkpoint has dummy dict
         self.assertTrue(
-            are_state_dict_equal(m1._module.state_dict(), m2._module.state_dict())
+            "foo" in loaded_checkpoint and loaded_checkpoint["foo"] == "bar"
         )
+        # check the two sets of components are now the same
+        self.assertEqual(pe1.accountant.state_dict(), pe2.accountant.state_dict())
+        self.assertTrue(are_state_dict_equal(m1.state_dict(), m2.state_dict()))
         if noise_scheduler:
             self.assertEqual(s1.state_dict(), s2.state_dict())
         # check that non-state params are still different
@@ -539,9 +544,7 @@ def test_checkpoints(self, noise_scheduler: Optional[_NoiseScheduler]):
             if noise_scheduler is not None
             else None
         )
-        self.assertFalse(
-            are_state_dict_equal(m2._module.state_dict(), m11._module.state_dict())
-        )
+        self.assertFalse(are_state_dict_equal(m2.state_dict(), m11.state_dict()))
         if noise_scheduler:
             self.assertNotEqual(s2.state_dict(), s11.state_dict())
         # train the recreated set for the same number of steps
@@ -552,9 +555,7 @@ def test_checkpoints(self, noise_scheduler: Optional[_NoiseScheduler]):
         if noise_scheduler:
             s11.step()
         # check that recreated set is now same as the original set 1 after training
-        self.assertTrue(
-            are_state_dict_equal(m2._module.state_dict(), m11._module.state_dict())
-        )
+        self.assertTrue(are_state_dict_equal(m2.state_dict(), m11.state_dict()))
         if noise_scheduler:
             self.assertEqual(s2.state_dict(), s11.state_dict())