Add enable_model_summary flag and deprecate weights_summary (#9699)

Co-authored-by: Adrian Wälchli <[email protected]>
Co-authored-by: Kaushik B <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kaushik B <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -181,6 +181,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for `torch.autograd.set_detect_anomaly` through `Trainer` constructor argument `detect_anomaly` ([#9848](https://github.com/PyTorchLightning/pytorch-lightning/pull/9848))
 
 
+- Added `enable_model_summary` flag to Trainer ([#9699](https://github.com/PyTorchLightning/pytorch-lightning/pull/9699))
+
+
 ### Changed
 
 - Module imports are now catching `ModuleNotFoundError` instead of `ImportError` ([#9867](https://github.com/PyTorchLightning/pytorch-lightning/pull/9867))
@@ -344,6 +347,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated `pytorch_lightning.core.decorators.parameter_validation` in favor of `pytorch_lightning.utilities.parameter_tying.set_shared_parameters` ([#9525](https://github.com/PyTorchLightning/pytorch-lightning/pull/9525))
 
 
+- Deprecated passing `weights_summary` to the `Trainer` constructor in favor of adding the `ModelSummary` callback with `max_depth` directly to the list of callbacks ([#9699](https://github.com/PyTorchLightning/pytorch-lightning/pull/9699))
+
+
 ### Removed
 
 - Removed deprecated `metrics` ([#8586](https://github.com/PyTorchLightning/pytorch-lightning/pull/8586/))

benchmarks/test_basic_parity.py

@@ -159,7 +159,7 @@ def lightning_loop(cls_model, idx, device_type: str = "cuda", num_epochs=10):
         # as the first run is skipped, no need to run it long
         max_epochs=num_epochs if idx > 0 else 1,
         enable_progress_bar=False,
-        weights_summary=None,
+        enable_model_summary=False,
         gpus=1 if device_type == "cuda" else 0,
         checkpoint_callback=False,
         logger=False,

docs/source/common/debugging.rst

@@ -95,11 +95,14 @@ Print a summary of your LightningModule
 ---------------------------------------
 Whenever the ``.fit()`` function gets called, the Trainer will print the weights summary for the LightningModule.
 By default it only prints the top-level modules. If you want to show all submodules in your network, use the
-`'full'` option:
+``max_depth`` option:
 
 .. testcode::
 
-    trainer = Trainer(weights_summary="full")
+    from pytorch_lightning.callbacks import ModelSummary
+
+    trainer = Trainer(callbacks=[ModelSummary(max_depth=-1)])
+
 
 You can also display the intermediate input- and output sizes of all your layers by setting the
 ``example_input_array`` attribute in your LightningModule. It will print a table like this
@@ -115,8 +118,9 @@ You can also display the intermediate input- and output sizes of all your layers
 when you call ``.fit()`` on the Trainer. This can help you find bugs in the composition of your layers.
 
 See Also:
-    - :paramref:`~pytorch_lightning.trainer.trainer.Trainer.weights_summary` Trainer argument
-    - :class:`~pytorch_lightning.core.memory.ModelSummary`
+    - :class:`~pytorch_lightning.callbacks.model_summary.ModelSummary`
+    - :func:`~pytorch_lightning.utilities.model_summary.summarize`
+    - :class:`~pytorch_lightning.utilities.model_summary.ModelSummary`
 
 ----------------
 

docs/source/common/trainer.rst

@@ -1589,6 +1589,11 @@ Example::
 weights_summary
 ^^^^^^^^^^^^^^^
 
+.. warning:: `weights_summary` is deprecated in v1.5 and will be removed in v1.7. Please pass :class:`~pytorch_lightning.callbacks.model_summary.ModelSummary`
+    directly to the Trainer's ``callbacks`` argument instead. To disable the model summary,
+    pass ``enable_model_summary = False`` to the Trainer.
+
+
 .. raw:: html
 
     <video width="50%" max-width="400px" controls
@@ -1611,6 +1616,25 @@ Options: 'full', 'top', None.
     # don't print a summary
     trainer = Trainer(weights_summary=None)
 
+
+enable_model_summary
+^^^^^^^^^^^^^^^^^^^^
+
+Whether to enable or disable the model summarization. Defaults to True.
+
+.. testcode::
+
+    # default used by the Trainer
+    trainer = Trainer(enable_model_summary=True)
+
+    # disable summarization
+    trainer = Trainer(enable_model_summary=False)
+
+    # enable custom summarization
+    from pytorch_lightning.callbacks import ModelSummary
+
+    trainer = Trainer(enable_model_summary=True, callbacks=[ModelSummary(max_depth=-1)])
+
 -----
 
 Trainer class API

pl_examples/bug_report_model.py

@@ -55,7 +55,7 @@ def run():
         limit_val_batches=1,
         num_sanity_val_steps=0,
         max_epochs=1,
-        weights_summary=None,
+        enable_model_summary=False,
     )
     trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
     trainer.test(model, dataloaders=test_data)

pl_examples/domain_templates/computer_vision_fine_tuning.py

@@ -272,7 +272,7 @@ def add_arguments_to_parser(self, parser):
         parser.set_defaults(
             {
                 "trainer.max_epochs": 15,
-                "trainer.weights_summary": None,
+                "trainer.enable_model_summary": False,
                 "trainer.num_sanity_val_steps": 0,
             }
         )

pytorch_lightning/trainer/connectors/callback_connector.py

@@ -45,6 +45,7 @@ def on_trainer_init(
         process_position: int,
         default_root_dir: Optional[str],
         weights_save_path: Optional[str],
+        enable_model_summary: bool,
         weights_summary: Optional[str],
         stochastic_weight_avg: bool,
         max_time: Optional[Union[str, timedelta, Dict[str, int]]] = None,
@@ -101,7 +102,7 @@ def on_trainer_init(
             self.trainer._progress_bar_callback = None
 
         # configure the ModelSummary callback
-        self._configure_model_summary_callback(weights_summary)
+        self._configure_model_summary_callback(enable_model_summary, weights_summary)
 
         # accumulated grads
         self._configure_accumulated_gradients(accumulate_grad_batches)
@@ -159,24 +160,50 @@ def _configure_checkpoint_callbacks(self, checkpoint_callback: Optional[bool], e
         if not self._trainer_has_checkpoint_callbacks() and enable_checkpointing is True:
             self.trainer.callbacks.append(ModelCheckpoint())
 
-    def _configure_model_summary_callback(self, weights_summary: Optional[str] = None) -> None:
-        if any(isinstance(cb, ModelSummary) for cb in self.trainer.callbacks):
+    def _configure_model_summary_callback(
+        self, enable_model_summary: bool, weights_summary: Optional[str] = None
+    ) -> None:
+        if weights_summary is None:
+            rank_zero_deprecation(
+                "Setting `Trainer(weights_summary=None)` is deprecated in v1.5 and will be removed"
+                " in v1.7. Please set `Trainer(enable_model_summary=False)` instead."
+            )
+            return
+        if not enable_model_summary:
+            return
+
+        model_summary_cbs = [type(cb) for cb in self.trainer.callbacks if isinstance(cb, ModelSummary)]
+        if model_summary_cbs:
+            rank_zero_info(
+                f"Trainer already configured with model summary callbacks: {model_summary_cbs}."
+                " Skipping setting a default `ModelSummary` callback."
+            )
             return
-        if weights_summary is not None:
+
+        if weights_summary == "top":
+            # special case the default value for weights_summary to preserve backward compatibility
+            max_depth = 1
+        else:
+            rank_zero_deprecation(
+                f"Setting `Trainer(weights_summary={weights_summary})` is deprecated in v1.5 and will be removed"
+                " in v1.7. Please pass `pytorch_lightning.callbacks.model_summary.ModelSummary` with"
+                " `max_depth` directly to the Trainer's `callbacks` argument instead."
+            )
             if weights_summary not in ModelSummaryMode.supported_types():
                 raise MisconfigurationException(
                     f"`weights_summary` can be None, {', '.join(ModelSummaryMode.supported_types())}",
                     f" but got {weights_summary}",
                 )
             max_depth = ModelSummaryMode.get_max_depth(weights_summary)
-            if self.trainer._progress_bar_callback is not None and isinstance(
-                self.trainer._progress_bar_callback, RichProgressBar
-            ):
-                model_summary = RichModelSummary(max_depth=max_depth)
-            else:
-                model_summary = ModelSummary(max_depth=max_depth)
-            self.trainer.callbacks.append(model_summary)
-            self.trainer.weights_summary = weights_summary
+
+        is_progress_bar_rich = isinstance(self.trainer._progress_bar_callback, RichProgressBar)
+
+        if self.trainer._progress_bar_callback is not None and is_progress_bar_rich:
+            model_summary = RichModelSummary(max_depth=max_depth)
+        else:
+            model_summary = ModelSummary(max_depth=max_depth)
+        self.trainer.callbacks.append(model_summary)
+        self.trainer._weights_summary = weights_summary
 
     def _configure_swa_callbacks(self):
         if not self.trainer._stochastic_weight_avg:

pytorch_lightning/trainer/trainer.py

@@ -157,6 +157,7 @@ def __init__(
         accelerator: Optional[Union[str, Accelerator]] = None,
         sync_batchnorm: bool = False,
         precision: Union[int, str] = 32,
+        enable_model_summary: bool = True,
         weights_summary: Optional[str] = "top",
         weights_save_path: Optional[str] = None,
         num_sanity_val_steps: int = 2,
@@ -373,8 +374,16 @@ def __init__(
             val_check_interval: How often to check the validation set. Use float to check within a training epoch,
                 use int to check every n steps (batches).
 
+            enable_model_summary: Whether to enable model summarization by default.
+
             weights_summary: Prints a summary of the weights when training begins.
 
+                .. deprecated:: v1.5
+                    ``weights_summary`` has been deprecated in v1.5 and will be removed in v1.7.
+                    To disable the summary, pass ``enable_model_summary = False`` to the Trainer.
+                    To customize the summary, pass :class:`~pytorch_lightning.callbacks.model_summary.ModelSummary`
+                    directly to the Trainer's ``callbacks`` argument.
+
             weights_save_path: Where to save weights if specified. Will override default_root_dir
                 for checkpoints only. Use this if for whatever reason you need the checkpoints
                 stored in a different place than the logs written in `default_root_dir`.
@@ -467,6 +476,9 @@ def __init__(
         self.tested_ckpt_path: Optional[str] = None
         self.predicted_ckpt_path: Optional[str] = None
 
+        # todo: remove in v1.7
+        self._weights_summary: Optional[str] = None
+
         # init callbacks
         # Declare attributes to be set in callback_connector on_trainer_init
         self.callback_connector.on_trainer_init(
@@ -478,6 +490,7 @@ def __init__(
             process_position,
             default_root_dir,
             weights_save_path,
+            enable_model_summary,
             weights_summary,
             stochastic_weight_avg,
             max_time,
@@ -2032,6 +2045,16 @@ def _exit_gracefully_on_signal(self) -> None:
             class_name = caller[0].f_locals["self"].__class__.__name__
             raise ExitGracefullyException(f"Exiting gracefully on {class_name}:{caller.function}")
 
+    @property
+    def weights_summary(self) -> Optional[str]:
+        rank_zero_deprecation("`Trainer.weights_summary` is deprecated in v1.5 and will be removed in v1.7.")
+        return self._weights_summary
+
+    @weights_summary.setter
+    def weights_summary(self, val: Optional[str]) -> None:
+        rank_zero_deprecation("Setting `Trainer.weights_summary` is deprecated in v1.5 and will be removed in v1.7.")
+        self._weights_summary = val
+
     """
     Other
     """

tests/accelerators/test_multi_nodes_gpu.py

@@ -54,7 +54,7 @@ def validation_step(self, batch, batch_idx):
         limit_train_batches=1,
         limit_val_batches=1,
         max_epochs=2,
-        weights_summary=None,
+        enable_model_summary=False,
         accelerator="ddp",
         gpus=1,
         num_nodes=2,
@@ -101,7 +101,7 @@ def backward(self, loss, optimizer, optimizer_idx):
         limit_val_batches=2,
         max_epochs=2,
         log_every_n_steps=1,
-        weights_summary=None,
+        enable_model_summary=False,
         accelerator="ddp",
         gpus=1,
         num_nodes=2,

tests/callbacks/test_callback_hook_outputs.py

@@ -53,7 +53,7 @@ def training_epoch_end(self, outputs) -> None:
         limit_val_batches=2,
         max_epochs=1,
         log_every_n_steps=1,
-        weights_summary=None,
+        enable_model_summary=False,
     )
 
     assert any(isinstance(c, CB) for c in trainer.callbacks)
@@ -74,7 +74,7 @@ def on_epoch_end(self, trainer, pl_module):
         limit_train_batches=2,
         limit_val_batches=2,
         max_epochs=1,
-        weights_summary=None,
+        enable_model_summary=False,
     )
 
     trainer.fit(model)