From 0dc0472e1f508a68035778474eb8e6d3b361fe5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Mochol=C3=AD?= Date: Thu, 29 Jul 2021 09:39:46 +0200 Subject: [PATCH] Use class name in SWA info message (#8602) --- pytorch_lightning/callbacks/stochastic_weight_avg.py | 5 ++++- tests/callbacks/test_stochastic_weight_avg.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/callbacks/stochastic_weight_avg.py b/pytorch_lightning/callbacks/stochastic_weight_avg.py index 219c2e7add..b681d2bff4 100644 --- a/pytorch_lightning/callbacks/stochastic_weight_avg.py +++ b/pytorch_lightning/callbacks/stochastic_weight_avg.py @@ -195,7 +195,10 @@ class StochasticWeightAveraging(Callback): scheduler_cfg = trainer.lr_schedulers[0] if scheduler_cfg["interval"] != "epoch" or scheduler_cfg["frequency"] != 1: rank_zero_warn(f"SWA is currently only supported every epoch. Found {scheduler_cfg}") - rank_zero_info(f"Swapping scheduler {scheduler_cfg['scheduler']} for {self._swa_scheduler}") + rank_zero_info( + f"Swapping scheduler `{scheduler_cfg['scheduler'].__class__.__name__}`" + f" for `{self._swa_scheduler.__class__.__name__}`" + ) trainer.lr_schedulers[0] = default_scheduler_cfg else: trainer.lr_schedulers.append(default_scheduler_cfg) diff --git a/tests/callbacks/test_stochastic_weight_avg.py b/tests/callbacks/test_stochastic_weight_avg.py index eed7ccd635..379a42e9a7 100644 --- a/tests/callbacks/test_stochastic_weight_avg.py +++ b/tests/callbacks/test_stochastic_weight_avg.py @@ -175,7 +175,7 @@ def test_swa_warns(tmpdir, caplog): trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, stochastic_weight_avg=True) with caplog.at_level(level=logging.INFO), pytest.warns(UserWarning, match="SWA is currently only supported"): trainer.fit(model) - assert "Swapping scheduler" in caplog.text + assert "Swapping scheduler `StepLR` for `SWALR`" in caplog.text def test_swa_raises():