trainer default to weights_only=None

Author: matsumotosan

src/lightning/pytorch/strategies/model_parallel.py

@@ -329,7 +329,7 @@ def save_checkpoint(
             return super().save_checkpoint(checkpoint=checkpoint, filepath=path)
 
     @override
-    def load_checkpoint(self, checkpoint_path: _PATH, weights_only: bool = False) -> dict[str, Any]:
+    def load_checkpoint(self, checkpoint_path: _PATH, weights_only: Optional[bool] = None) -> dict[str, Any]:
         # broadcast the path from rank 0 to ensure all the states are loaded from a common path
         path = Path(self.broadcast(checkpoint_path))
         state = {

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -64,7 +64,7 @@ def _hpc_resume_path(self) -> Optional[str]:
             return dir_path_hpc + fs.sep + f"hpc_ckpt_{max_version}.ckpt"
         return None
 
-    def resume_start(self, checkpoint_path: Optional[_PATH] = None, weights_only: bool = False) -> None:
+    def resume_start(self, checkpoint_path: Optional[_PATH] = None, weights_only: Optional[bool] = None) -> None:
         """Attempts to pre-load the checkpoint file to memory, with the source path determined in this priority:
 
         1. from HPC weights if `checkpoint_path` is ``None`` and on SLURM or passed keyword `"hpc"`.
@@ -404,7 +404,7 @@ def restore_lr_schedulers(self) -> None:
             config.scheduler.load_state_dict(lrs_state)
 
     def _restore_modules_and_callbacks(
-        self, checkpoint_path: Optional[_PATH] = None, weights_only: bool = False
+        self, checkpoint_path: Optional[_PATH] = None, weights_only: Optional[bool] = None
     ) -> None:
         # restore modules after setup
         self.resume_start(checkpoint_path, weights_only)

src/lightning/pytorch/trainer/trainer.py

@@ -526,7 +526,7 @@ def fit(
         val_dataloaders: Optional[EVAL_DATALOADERS] = None,
         datamodule: Optional[LightningDataModule] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
     ) -> None:
         r"""Runs the full optimization routine.
 
@@ -591,7 +591,7 @@ def _fit_impl(
         val_dataloaders: Optional[EVAL_DATALOADERS] = None,
         datamodule: Optional[LightningDataModule] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
     ) -> None:
         log.debug(f"{self.__class__.__name__}: trainer fit stage")
 
@@ -630,7 +630,7 @@ def validate(
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
     ) -> _EVALUATE_OUTPUT:
@@ -694,7 +694,7 @@ def _validate_impl(
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
     ) -> Optional[Union[_PREDICT_OUTPUT, _EVALUATE_OUTPUT]]:
@@ -742,7 +742,7 @@ def test(
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
     ) -> _EVALUATE_OUTPUT:
@@ -807,7 +807,7 @@ def _test_impl(
         model: Optional["pl.LightningModule"] = None,
         dataloaders: Optional[Union[EVAL_DATALOADERS, LightningDataModule]] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
         verbose: bool = True,
         datamodule: Optional[LightningDataModule] = None,
     ) -> Optional[Union[_PREDICT_OUTPUT, _EVALUATE_OUTPUT]]:
@@ -857,7 +857,7 @@ def predict(
         datamodule: Optional[LightningDataModule] = None,
         return_predictions: Optional[bool] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
     ) -> Optional[_PREDICT_OUTPUT]:
         r"""Run inference on your data. This will call the model forward function to compute predictions. Useful to
         perform distributed and batched predictions. Logging is disabled in the predict hooks.
@@ -923,7 +923,7 @@ def _predict_impl(
         datamodule: Optional[LightningDataModule] = None,
         return_predictions: Optional[bool] = None,
         ckpt_path: Optional[_PATH] = None,
-        weights_only: bool = False,
+        weights_only: Optional[bool] = None,
     ) -> Optional[_PREDICT_OUTPUT]:
         # --------------------
         # SETUP HOOK
@@ -962,7 +962,10 @@ def _predict_impl(
         return results
 
     def _run(
-        self, model: "pl.LightningModule", ckpt_path: Optional[_PATH] = None, weights_only: bool = False
+        self,
+        model: "pl.LightningModule",
+        ckpt_path: Optional[_PATH] = None,
+        weights_only: Optional[bool] = None,
     ) -> Optional[Union[_EVALUATE_OUTPUT, _PREDICT_OUTPUT]]:
         if self.state.fn == TrainerFn.FITTING:
             min_epochs, max_epochs = _parse_loop_limits(
@@ -1401,7 +1404,7 @@ def ckpt_path(self, ckpt_path: Optional[_PATH]) -> None:
         self._checkpoint_connector._user_managed = bool(ckpt_path)
 
     def save_checkpoint(
-        self, filepath: _PATH, weights_only: bool = False, storage_options: Optional[Any] = None
+        self, filepath: _PATH, weights_only: Optional[bool] = None, storage_options: Optional[Any] = None
     ) -> None:
         r"""Runs routine to create a checkpoint.