Update tests to avoid the deprecated weights_summary

pytorch_lightning/core/lightning.py

@@ -1708,7 +1708,7 @@ def summarize(self, mode: Optional[str] = "top", max_depth: Optional[int] = None
         Return:
             The model summary object
         """
-        warning_cache.deprecation(
+        rank_zero_deprecation(
             "The `LightningModule.summarize` method is deprecated in v1.5 and will be removed in v1.7. "
             "Use `pytorch_lightning.utilities.model_summary.summarize` instead.",
             stacklevel=6,

pytorch_lightning/tuner/batch_size_scaling.py

@@ -106,7 +106,6 @@ def __scale_batch_dump_params(trainer: "pl.Trainer") -> None:
         "current_epoch": trainer.current_epoch,
         "global_step": trainer.global_step,
         "max_steps": trainer.max_steps,
-        "weights_summary": trainer.weights_summary,
         "logger": trainer.logger,
         "callbacks": trainer.callbacks,
         "checkpoint_callback": trainer.checkpoint_callback,
@@ -121,7 +120,6 @@ def __scale_batch_reset_params(trainer: "pl.Trainer", model: "pl.LightningModule
     trainer.auto_lr_find = False  # avoid lr find being called multiple times
     trainer.fit_loop.current_epoch = 0
     trainer.fit_loop.max_steps = steps_per_trial  # take few steps
-    trainer.weights_summary = None  # not needed before full run
     trainer.logger = DummyLogger() if trainer.logger is not None else None
     trainer.callbacks = []  # not needed before full run
     trainer.limit_train_batches = 1.0
@@ -134,7 +132,6 @@ def __scale_batch_restore_params(trainer: "pl.Trainer") -> None:
     trainer.fit_loop.current_epoch = trainer.__dumped_params["current_epoch"]
     trainer.fit_loop.global_step = trainer.__dumped_params["global_step"]
     trainer.fit_loop.max_steps = trainer.__dumped_params["max_steps"]
-    trainer.weights_summary = trainer.__dumped_params["weights_summary"]
     trainer.logger = trainer.__dumped_params["logger"]
     trainer.callbacks = trainer.__dumped_params["callbacks"]
     trainer.auto_scale_batch_size = trainer.__dumped_params["auto_scale_batch_size"]

tests/callbacks/test_model_summary.py

@@ -17,8 +17,6 @@
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelSummary
-from pytorch_lightning.utilities import ModelSummaryMode
-from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
 
 
@@ -48,26 +46,19 @@ def test_model_summary_callback_with_weights_summary_none():
 
 
 def test_model_summary_callback_with_weights_summary():
-
     trainer = Trainer(weights_summary="top")
-
     model_summary_callback = list(filter(lambda cb: isinstance(cb, ModelSummary), trainer.callbacks))[0]
     assert model_summary_callback._max_depth == 1
 
-    trainer = Trainer(weights_summary="full")
-
+    with pytest.deprecated_call(match=r"weights_summary=full\)` is deprecated"):
+        trainer = Trainer(weights_summary="full")
     model_summary_callback = list(filter(lambda cb: isinstance(cb, ModelSummary), trainer.callbacks))[0]
     assert model_summary_callback._max_depth == -1
 
-    with pytest.raises(
-        MisconfigurationException, match=f"`weights_summary` can be None, {', '.join(list(ModelSummaryMode))}"
-    ):
-        _ = Trainer(weights_summary="invalid")
-
 
 def test_model_summary_callback_override_weights_summary_flag():
-
-    trainer = Trainer(callbacks=ModelSummary(), weights_summary=None)
+    with pytest.deprecated_call(match=r"weights_summary=None\)` is deprecated"):
+        trainer = Trainer(callbacks=ModelSummary(), weights_summary=None)
     assert any(isinstance(cb, ModelSummary) for cb in trainer.callbacks)
 
 

tests/deprecated_api/test_remove_1-7.py

@@ -38,12 +38,9 @@
 
 
 def test_v1_7_0_deprecated_lightning_module_summarize(tmpdir):
-    from pytorch_lightning.core.lightning import warning_cache
-
     model = BoringModel()
-    model.summarize(max_depth=1)
-    assert any("The `LightningModule.summarize` method is deprecated in v1.5" in w for w in warning_cache)
-    warning_cache.clear()
+    with pytest.deprecated_call(match="The `LightningModule.summarize` method is deprecated in v1.5"):
+        model.summarize(max_depth=1)
 
 
 def test_v1_7_0_moved_model_summary_and_layer_summary(tmpdir):

tests/tuner/test_scale_batch_size.py

@@ -114,7 +114,6 @@ def test_trainer_reset_correctly(tmpdir):
         "logger",
         "max_steps",
         "global_step",
-        "weights_summary",
     ]
     expected = {ca: getattr(trainer, ca) for ca in changed_attributes}
     trainer.tuner.scale_batch_size(model, max_trials=5)

tests/utilities/test_model_summary.py

@@ -141,31 +141,41 @@ def forward(self, inp):
 
 def test_invalid_weights_summmary():
     """Test that invalid value for weights_summary raises an error."""
+    model = LightningModule()
+
     with pytest.raises(MisconfigurationException, match="`mode` can be None, .* got temp"):
-        summarize(UnorderedModel, mode="temp")
+        summarize(model, mode="temp")
 
-    with pytest.raises(MisconfigurationException, match="`weights_summary` can be None, .* got temp"):
+    with pytest.raises(
+        MisconfigurationException, match="`weights_summary` can be None, .* got temp"
+    ), pytest.deprecated_call(match="weights_summary=temp)` is deprecated"):
         Trainer(weights_summary="temp")
 
+    with pytest.raises(MisconfigurationException, match="mode` can be .* got temp"):
+        ModelSummary(model, mode="temp")
+
+    with pytest.raises(ValueError, match="max_depth` can be .* got temp"):
+        ModelSummary(model, max_depth="temp")
+
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_empty_model_summary_shapes(mode: str):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_empty_model_summary_shapes(max_depth):
     """Test that the summary works for models that have no submodules."""
     model = EmptyModule()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.in_sizes == []
     assert summary.out_sizes == []
     assert summary.param_nums == []
 
 
 @RunIf(min_gpus=1)
-@pytest.mark.parametrize("mode", ["full", "top"])
+@pytest.mark.parametrize("max_depth", [-1, 1])
 @pytest.mark.parametrize("device", [torch.device("cpu"), torch.device("cuda", 0)])
-def test_linear_model_summary_shapes(device, mode):
+def test_linear_model_summary_shapes(device, max_depth):
     """Test that the model summary correctly computes the input- and output shapes."""
     model = UnorderedModel().to(device)
     model.train()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.in_sizes == [[2, 10], [2, 7], [2, 3], [2, 7], UNKNOWN_SIZE]  # layer 2  # combine  # layer 1  # relu
     assert summary.out_sizes == [[2, 2], [2, 9], [2, 5], [2, 7], UNKNOWN_SIZE]  # layer 2  # combine  # layer 1  # relu
     assert model.training
@@ -191,8 +201,8 @@ def test_hooks_removed_after_summarize(max_depth):
         assert handle.id not in handle.hooks_dict_ref()
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_rnn_summary_shapes(mode):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_rnn_summary_shapes(max_depth):
     """Test that the model summary works for RNNs."""
     model = ParityModuleRNN()
 
@@ -204,16 +214,16 @@ def test_rnn_summary_shapes(mode):
 
     model.example_input_array = torch.zeros(b, t, 10)
 
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.in_sizes == [[b, t, i], [b, t, h]]  # rnn  # linear
     assert summary.out_sizes == [[[b, t, h], [[1, b, h], [1, b, h]]], [b, t, o]]  # rnn  # linear
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_summary_parameter_count(mode):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_summary_parameter_count(max_depth):
     """Test that the summary counts the number of parameters in every submodule."""
     model = UnorderedModel()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.param_nums == [
         model.layer2.weight.numel() + model.layer2.bias.numel(),
         model.combine.weight.numel() + model.combine.bias.numel(),
@@ -223,24 +233,24 @@ def test_summary_parameter_count(mode):
     ]
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_summary_layer_types(mode):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_summary_layer_types(max_depth):
     """Test that the summary displays the layer names correctly."""
     model = UnorderedModel()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.layer_types == ["Linear", "Linear", "Linear", "ReLU", "Conv2d"]
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_summary_with_scripted_modules(mode):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_summary_with_scripted_modules(max_depth):
     model = PartialScriptModel()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.layer_types == ["RecursiveScriptModule", "Linear"]
     assert summary.in_sizes == [UNKNOWN_SIZE, [2, 3]]
     assert summary.out_sizes == [UNKNOWN_SIZE, [2, 2]]
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
+@pytest.mark.parametrize("max_depth", [-1, 1])
 @pytest.mark.parametrize(
     ["example_input", "expected_size"],
     [
@@ -253,7 +263,7 @@ def test_summary_with_scripted_modules(mode):
         ((torch.zeros(2, 3), torch.zeros(4, 5)), [[2, 3], [4, 5]]),
     ],
 )
-def test_example_input_array_types(example_input, expected_size, mode):
+def test_example_input_array_types(example_input, expected_size, max_depth):
     """Test the types of example inputs supported for display in the summary."""
 
     class DummyModule(nn.Module):
@@ -271,23 +281,23 @@ def forward(self, *args, **kwargs):
 
     model = DummyLightningModule()
     model.example_input_array = example_input
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert summary.in_sizes == [expected_size]
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_model_size(mode):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_model_size(max_depth):
     """Test model size is calculated correctly."""
     model = PreCalculatedModel()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert model.pre_calculated_model_size == summary.model_size
 
 
-@pytest.mark.parametrize("mode", ["full", "top"])
-def test_empty_model_size(mode):
+@pytest.mark.parametrize("max_depth", [-1, 1])
+def test_empty_model_size(max_depth):
     """Test empty model size is zero."""
     model = EmptyModule()
-    summary = summarize(model, mode=mode)
+    summary = summarize(model, max_depth=max_depth)
     assert 0.0 == summary.model_size
 
 
@@ -328,11 +338,13 @@ def test_max_depth_equals_mode_interface():
     """Test summarize(model, full/top) interface mapping matches max_depth."""
     model = DeepNestedModel()
 
-    summary_top = summarize(model, mode="top")
+    with pytest.deprecated_call(match="mode` in `LightningModule.summarize` is deprecated"):
+        summary_top = summarize(model, mode="top")
     summary_0 = summarize(model, max_depth=1)
     assert str(summary_top) == str(summary_0)
 
-    summary_full = summarize(model, mode="full")
+    with pytest.deprecated_call(match="mode` in `LightningModule.summarize` is deprecated"):
+        summary_full = summarize(model, mode="full")
     summary_minus1 = summarize(model, max_depth=-1)
     assert str(summary_full) == str(summary_minus1)