Deprecate TrainingTypePlugin.on_save and Accelerator.on_save

Author: four4fish

CHANGELOG.md

@@ -111,7 +111,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated `DataModule` properties: `train_transforms`, `val_transforms`, `test_transforms`, `size`, `dims` ([#8851](https://github.com/PyTorchLightning/pytorch-lightning/pull/8851))
 
 
--
+- Deprecated `TrainingTypePlugin.on_save` and `Accelerator.on_save` ([#8987](https://github.com/PyTorchLightning/pytorch-lightning/pull/8987))
 
 
 -

pytorch_lightning/accelerators/accelerator.py

@@ -371,9 +371,6 @@ def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]:
         """
         return self.training_type_plugin.lightning_module_state_dict()
 
-    def on_save(self, checkpoint: Dict[str, Union[Any, Tensor]]) -> Dict[str, Union[Any, Tensor]]:
-        return self.training_type_plugin.on_save(checkpoint)
-
     def barrier(self, name: Optional[str] = None) -> None:
         self.training_type_plugin.barrier(name=name)
 

pytorch_lightning/plugins/training_type/ddp_spawn.py

@@ -301,7 +301,7 @@ def __transfer_distrib_spawn_state_on_fit_end(self, trainer: "pl.Trainer", resul
             last_path = None
             if trainer.state.fn == TrainerFn.FITTING and best_model_path is not None and len(best_model_path) > 0:
                 last_path = re.sub(".ckpt", ".tmp_end.ckpt", best_model_path)
-                atomic_save(self.on_save(state_dict), last_path)
+                atomic_save(state_dict, last_path)
 
             # todo, pass complete checkpoint as state dictionary
             self.mp_queue.put(best_model_path)

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -201,9 +201,6 @@ def validation_step_end(self, output):
     def test_step_end(self, output):
         return output
 
-    def on_save(self, checkpoint: Dict[str, Union[Any, torch.Tensor]]) -> Dict[str, Union[Any, torch.Tensor]]:
-        return checkpoint
-
     def process_dataloader(self, dataloader: Union[Iterable, DataLoader]) -> Union[Iterable, DataLoader]:
         """Wraps the dataloader if necessary
 
@@ -273,8 +270,6 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], filepath: str) -> None:
             checkpoint: dict containing model and trainer state
             filepath: write-target file's path
         """
-        # dump states as a checkpoint dictionary object
-        checkpoint = self.on_save(checkpoint)
         if self.should_rank_save_checkpoint:
             return self.checkpoint_io.save_checkpoint(checkpoint, filepath)
 

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -294,8 +294,6 @@ def hpc_save(self, folderpath: str, logger):
 
         model.on_hpc_save(checkpoint)
 
-        checkpoint = self.trainer.accelerator.on_save(checkpoint)
-
         # do the actual save
         # TODO: fix for anything with multiprocess DP, DDP, DDP2
         try: