Repository: pytorch/ignite

Issue:
WeightsHistHandler should plot all weights (incl those without grad)
## 🚀 Feature

Currently, `WeightsHistHandler` do not log weights without grad:
https://github.com/pytorch/ignite/blob/a168476762061f74af64793bce1b2d5b6e9b3f15/ignite/contrib/handlers/tensorboard_logger.py#L434-L435

Let's put an option to enable logging all weights. Maybe, we can make it even without option and just log everything ?


Meanwhile, here is a workaround code for TensorboardLogger:
```python
from ignite.contrib.handlers import TensorboardLogger
from ignite.contrib.handlers.tensorboard_logger import WeightsHistHandler

class FixedWeightsHistHandler(WeightsHistHandler):
    
    def __call__(self, engine, logger, event_name):
        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():            
            name = name.replace(".", "/")
            logger.writer.add_histogram(
                tag=f"{tag_prefix}weights/{name}", values=p.data.detach().cpu().numpy(), global_step=global_step,
            )

```

WeightsHistHandler should plot all weights (incl those without grad)
## 🚀 Feature

Currently, `WeightsHistHandler` do not log weights without grad:
https://github.com/pytorch/ignite/blob/a168476762061f74af64793bce1b2d5b6e9b3f15/ignite/contrib/handlers/tensorboard_logger.py#L434-L435

Let's put an option to enable logging all weights. Maybe, we can make it even without option and just log everything ?


Meanwhile, here is a workaround code for TensorboardLogger:
```python
from ignite.contrib.handlers import TensorboardLogger
from ignite.contrib.handlers.tensorboard_logger import WeightsHistHandler

class FixedWeightsHistHandler(WeightsHistHandler):
    
    def __call__(self, engine, logger, event_name):
        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():            
            name = name.replace(".", "/")
            logger.writer.add_histogram(
                tag=f"{tag_prefix}weights/{name}", values=p.data.detach().cpu().numpy(), global_step=global_step,
            )

```



Discussion:
Assume you have a partial frozen model, it makes sense to only trace the evolution of the weight which are updated.
The question is does checking if p.grad is not None is what we want ?
What about weights modification with other computation ? like pca_low_rank ?

I see what you mean. Ultimately, we would like to have an optional list of weights to plot ? 
Something like:

```python
from ignite.contrib.handlers import TensorboardLogger
from ignite.contrib.handlers.tensorboard_logger import WeightsHistHandler


class FixedWeightsHistHandler(WeightsHistHandler):

    def __init__(self, model, tag=None, names_whitelist=None):
        super(WeightsHistHandler, self).__init__(model, tag=tag)
        self.names_whitelist = names_whitelist
    
    def __call__(self, engine, logger, event_name):
        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 self.names_whitelist is not None and name not in self.names_whitelist:
                continue

            name = name.replace(".", "/")
            logger.writer.add_histogram(
                tag=f"{tag_prefix}weights/{name}", values=p.data.detach().cpu().numpy(), global_step=global_step,
            )


tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")
# ...
tb_logger.attach(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    log_handler=FixedWeightsHistHandler(model, names_whitelist=[n for n, _ in model.named_parameters() if "conv" in n])
)
```
@DhDeepLIT I thought you wanted to avoid the case of inexisting grads and log weights in any case ?
```python
if p.grad is None:
    continue
```
Or you are now talking about `GradsHistHandler` ?
Yes that is my need, but on your side, what was THE thing you wanted to do ?
Consider that the zero_grad() function can also set grad to None.
Thus checking if the grad attribute is None or Not may not be the right option.
BTW your proposed patch is fine for my need, the whitelist is a great idea too.
FYI: zero_grad does not set grads to None by default but just zeros them :)

https://github.com/pytorch/pytorch/blob/33e9a0b5f6674bd429bda689534e1c987b38cf6e/torch/nn/modules/module.py#L1785-L1794

Maybe it should be the other way around though :
it is easier to filter the weights/grad we want if a key is present in the module full name say you want to check ALL fc layer. It would be better to check if 'fc' is in the named parameter of the module, right ?
I agree that non-exact match for names_whitelist could be helpful...
https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html

There is a option with 'set_to_none=True'
Assume you have a partial frozen model, it makes sense to only trace the evolution of the weight which are updated.
The question is does checking if p.grad is not None is what we want ?
What about weights modification with other computation ? like pca_low_rank ?

I see what you mean. Ultimately, we would like to have an optional list of weights to plot ? 
Something like:

```python
from ignite.contrib.handlers import TensorboardLogger
from ignite.contrib.handlers.tensorboard_logger import WeightsHistHandler


class FixedWeightsHistHandler(WeightsHistHandler):

    def __init__(self, model, tag=None, names_whitelist=None):
        super(WeightsHistHandler, self).__init__(model, tag=tag)
        self.names_whitelist = names_whitelist
    
    def __call__(self, engine, logger, event_name):
        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 self.names_whitelist is not None and name not in self.names_whitelist:
                continue

            name = name.replace(".", "/")
            logger.writer.add_histogram(
                tag=f"{tag_prefix}weights/{name}", values=p.data.detach().cpu().numpy(), global_step=global_step,
            )


tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")
# ...
tb_logger.attach(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    log_handler=FixedWeightsHistHandler(model, names_whitelist=[n for n, _ in model.named_parameters() if "conv" in n])
)
```
@DhDeepLIT I thought you wanted to avoid the case of inexisting grads and log weights in any case ?
```python
if p.grad is None:
    continue
```
Or you are now talking about `GradsHistHandler` ?
Yes that is my need, but on your side, what was THE thing you wanted to do ?
Consider that the zero_grad() function can also set grad to None.
Thus checking if the grad attribute is None or Not may not be the right option.
BTW your proposed patch is fine for my need, the whitelist is a great idea too.
FYI: zero_grad does not set grads to None by default but just zeros them :)

https://github.com/pytorch/pytorch/blob/33e9a0b5f6674bd429bda689534e1c987b38cf6e/torch/nn/modules/module.py#L1785-L1794

Maybe it should be the other way around though :
it is easier to filter the weights/grad we want if a key is present in the module full name say you want to check ALL fc layer. It would be better to check if 'fc' is in the named parameter of the module, right ?
I agree that non-exact match for names_whitelist could be helpful...
https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html

There is a option with 'set_to_none=True'

Resolving patch:
diff --git a/examples/contrib/mnist/mnist_with_tensorboard_logger.py b/examples/contrib/mnist/mnist_with_tensorboard_logger.py
--- a/examples/contrib/mnist/mnist_with_tensorboard_logger.py
+++ b/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"]
diff --git a/ignite/contrib/handlers/tensorboard_logger.py b/ignite/contrib/handlers/tensorboard_logger.py
--- a/ignite/contrib/handlers/tensorboard_logger.py
+++ b/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(


Test patch:
diff --git a/tests/ignite/contrib/handlers/conftest.py b/tests/ignite/contrib/handlers/conftest.py
--- a/tests/ignite/contrib/handlers/conftest.py
+++ b/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
diff --git a/tests/ignite/contrib/handlers/test_tensorboard_logger.py b/tests/ignite/contrib/handlers/test_tensorboard_logger.py
--- a/tests/ignite/contrib/handlers/test_tensorboard_logger.py
+++ b/tests/ignite/contrib/handlers/test_tensorboard_logger.py
@@ -380,11 +380,10 @@ 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()
 
@@ -392,24 +391,36 @@ def test_weights_hist_handler_frozen_layers(dummy_model_factory):
     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