Feature fixed weights hist handler issue 2328 (#2523)

* Remove unnecessary code in BaseOutputHandler

Closes #2438

* Add ReduceLROnPlateauScheduler

Closes #1754

* Fix indentation issue

* Fix another indentation issue

* Fix PEP8 related issues

* Fix other PEP8 related issues

* Fix hopefully the last PEP8 related issue

* Fix hopefully the last PEP8 related issue

* Remove ReduceLROnPlateau's specific params and add link to it

Also fix bug in min_lr check

* Fix state_dict bug and add a test

* Update docs

* Add doctest and fix typo

* Add feature FixedWeightsHistHandler

Closes #2328

* Move FixedWeightsHistHandler's job to WeightsHistHandler

Closes #2328

* Enable whitelist to be callable

* autopep8 fix

* Refactor constructor

* Change whitelist to be List[str]

* Add whitelist callable type

* Fix bug in MNIST tensorboard example

Co-authored-by: vfdev <[email protected]>
Co-authored-by: sadra-barikbin <[email protected]>

Author: sadra-barikbin

examples/contrib/mnist/mnist_with_tensorboard_logger.py

@@ -135,11 +135,20 @@ def compute_metrics(engine):
 
     tb_logger.attach(trainer, log_handler=WeightsScalarHandler(model), event_name=Events.ITERATION_COMPLETED(every=100))
 
-    tb_logger.attach(trainer, log_handler=WeightsHistHandler(model), event_name=Events.EPOCH_COMPLETED(every=100))
+    tb_logger.attach(
+        trainer,
+        log_handler=WeightsHistHandler(
+            model,
+            whitelist=[
+                "conv",
+            ],
+        ),
+        event_name=Events.ITERATION_COMPLETED(every=100),
+    )
 
     tb_logger.attach(trainer, log_handler=GradsScalarHandler(model), event_name=Events.ITERATION_COMPLETED(every=100))
 
-    tb_logger.attach(trainer, log_handler=GradsHistHandler(model), event_name=Events.EPOCH_COMPLETED(every=100))
+    tb_logger.attach(trainer, log_handler=GradsHistHandler(model), event_name=Events.ITERATION_COMPLETED(every=100))
 
     def score_function(engine):
         return engine.state.metrics["accuracy"]

ignite/contrib/handlers/tensorboard_logger.py

@@ -404,6 +404,12 @@ class WeightsHistHandler(BaseWeightsHistHandler):
     Args:
         model: model to log weights
         tag: common title for all produced plots. For example, "generator"
+        whitelist: specific weights to log. Should be list of model's submodules
+            or parameters names, or a callable which gets weight along with its name
+            and determines if it should be logged. Names should be fully-qualified.
+            For more information please refer to `PyTorch docs
+            <https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.get_submodule>`_.
+            If not given, all of model's weights are logged.
 
     Examples:
         .. code-block:: python
@@ -419,20 +425,75 @@ class WeightsHistHandler(BaseWeightsHistHandler):
                 event_name=Events.ITERATION_COMPLETED,
                 log_handler=WeightsHistHandler(model)
             )
+
+        .. code-block:: python
+
+            from ignite.contrib.handlers.tensorboard_logger import *
+
+            # Create a logger
+            tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")
+
+            # Log weights of `fc` layer
+            weights = ['fc']
+
+            # Attach the logger to the trainer to log weights norm after each iteration
+            tb_logger.attach(
+                trainer,
+                event_name=Events.ITERATION_COMPLETED,
+                log_handler=WeightsHistHandler(model, whitelist=weights)
+            )
+
+        .. code-block:: python
+
+            from ignite.contrib.handlers.tensorboard_logger import *
+
+            # Create a logger
+            tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")
+
+            # Log weights which name include 'conv'.
+            weight_selector = lambda name, p: 'conv' in name
+
+            # Attach the logger to the trainer to log weights norm after each iteration
+            tb_logger.attach(
+                trainer,
+                event_name=Events.ITERATION_COMPLETED,
+                log_handler=WeightsHistHandler(model, whitelist=weight_selector)
+            )
     """
 
-    def __init__(self, model: nn.Module, tag: Optional[str] = None):
+    def __init__(
+        self,
+        model: nn.Module,
+        tag: Optional[str] = None,
+        whitelist: Optional[Union[List[str], Callable[[str, nn.Parameter], bool]]] = None,
+    ):
         super(WeightsHistHandler, self).__init__(model, tag=tag)
 
+        weights = {}
+        if whitelist is None:
+
+            weights = dict(model.named_parameters())
+        elif callable(whitelist):
+
+            for n, p in model.named_parameters():
+                if whitelist(n, p):
+                    weights[n] = p
+        else:
+
+            for n, p in model.named_parameters():
+                for item in whitelist:
+                    if n.startswith(item):
+                        weights[n] = p
+
+        self.weights = weights.items()
+
     def __call__(self, engine: Engine, logger: TensorboardLogger, event_name: Union[str, Events]) -> None:
         if not isinstance(logger, TensorboardLogger):
             raise RuntimeError("Handler 'WeightsHistHandler' works only with TensorboardLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
         tag_prefix = f"{self.tag}/" if self.tag else ""
-        for name, p in self.model.named_parameters():
-            if p.grad is None:
-                continue
+        for name, p in self.weights:
 
             name = name.replace(".", "/")
             logger.writer.add_histogram(

tests/ignite/contrib/handlers/conftest.py

@@ -28,7 +28,7 @@ def __init__(self):
             self.fc2.weight.data.fill_(1.0)
             self.fc2.bias.data.fill_(1.0)
 
-    def get_dummy_model(with_grads=True, with_frozen_layer=False):
+    def get_dummy_model(with_grads=True, with_frozen_layer=False, with_buffer=False):
         model = DummyModel()
         if with_grads:
             model.fc2.weight.grad = torch.zeros_like(model.fc2.weight)
@@ -41,6 +41,9 @@ def get_dummy_model(with_grads=True, with_frozen_layer=False):
         if with_frozen_layer:
             for param in model.fc1.parameters():
                 param.requires_grad = False
+
+        if with_buffer:
+            model.register_buffer("buffer1", torch.ones(1))
         return model
 
     return get_dummy_model

tests/ignite/contrib/handlers/test_tensorboard_logger.py

@@ -380,36 +380,47 @@ def _test(tag=None):
     _test(tag="tag")
 
 
-def test_weights_hist_handler_frozen_layers(dummy_model_factory):
+def test_weights_hist_handler_whitelist(dummy_model_factory):
+    model = dummy_model_factory()
 
-    model = dummy_model_factory(with_grads=True, with_frozen_layer=True)
-
-    wrapper = WeightsHistHandler(model)
+    wrapper = WeightsHistHandler(model, whitelist=["fc2.weight"])
     mock_logger = MagicMock(spec=TensorboardLogger)
     mock_logger.writer = MagicMock()
 
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 5
 
+    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+    mock_logger.writer.add_histogram.assert_called_once_with(tag="weights/fc2/weight", values=ANY, global_step=5)
+    mock_logger.writer.reset_mock()
+
+    wrapper = WeightsHistHandler(model, tag="model", whitelist=["fc1"])
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
     mock_logger.writer.add_histogram.assert_has_calls(
         [
-            call(tag="weights/fc2/weight", values=ANY, global_step=5),
-            call(tag="weights/fc2/bias", values=ANY, global_step=5),
+            call(tag="model/weights/fc1/weight", values=ANY, global_step=5),
+            call(tag="model/weights/fc1/bias", values=ANY, global_step=5),
         ],
         any_order=True,
     )
+    assert mock_logger.writer.add_histogram.call_count == 2
+    mock_logger.writer.reset_mock()
 
-    with pytest.raises(AssertionError):
-        mock_logger.writer.add_histogram.assert_has_calls(
-            [
-                call(tag="weights/fc1/weight", values=ANY, global_step=5),
-                call(tag="weights/fc1/bias", values=ANY, global_step=5),
-            ],
-            any_order=True,
-        )
+    def weight_selector(n, _):
+        return "bias" in n
+
+    wrapper = WeightsHistHandler(model, tag="model", whitelist=weight_selector)
+    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+
+    mock_logger.writer.add_histogram.assert_has_calls(
+        [
+            call(tag="model/weights/fc1/bias", values=ANY, global_step=5),
+            call(tag="model/weights/fc2/bias", values=ANY, global_step=5),
+        ],
+        any_order=True,
+    )
     assert mock_logger.writer.add_histogram.call_count == 2