Repository: pytorch/ignite

Issue:
Tensorboard Logger does not work with layer freezing
Consider the following (simplified and self-contained) notebook:
https://gist.github.com/marrrcin/629227cf33d88af6192048871daa3a09
which uses mostly the code from Ignite examples to create `trainer` and attach `TensorboardLogger` to training events.

If you use pre-trained model and freeze some layers, you end up with code that will not work, as Ignite tries to collect gradient histograms/scalars for every model parameter, even though some of them are already disabled from updating.

The exception thrown is the following:

```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-3-5ffffec84f3c> in <module>
     65     print("Done.")
     66 
---> 67 example()

<ipython-input-3-5ffffec84f3c> in example()
     62 
     63 
---> 64     trainer.run(DataLoader(FakeDataset(), batch_size=16), max_epochs=5)
     65     print("Done.")
     66 

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in run(self, data, max_epochs)
    357         except BaseException as e:
    358             self._logger.error("Engine run is terminating due to exception: %s.", str(e))
--> 359             self._handle_exception(e)
    360 
    361         return self.state

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in _handle_exception(self, e)
    322             self._fire_event(Events.EXCEPTION_RAISED, e)
    323         else:
--> 324             raise e
    325 
    326     def run(self, data, max_epochs=1):

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in run(self, data, max_epochs)
    344                 self.state.epoch += 1
    345                 self._fire_event(Events.EPOCH_STARTED)
--> 346                 hours, mins, secs = self._run_once_on_dataset()
    347                 self._logger.info("Epoch[%s] Complete. Time taken: %02d:%02d:%02d", self.state.epoch, hours, mins, secs)
    348                 if self.should_terminate:

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in _run_once_on_dataset(self)
    311         except BaseException as e:
    312             self._logger.error("Current run is terminating due to exception: %s.", str(e))
--> 313             self._handle_exception(e)
    314 
    315         time_taken = time.time() - start_time

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in _handle_exception(self, e)
    322             self._fire_event(Events.EXCEPTION_RAISED, e)
    323         else:
--> 324             raise e
    325 
    326     def run(self, data, max_epochs=1):

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in _run_once_on_dataset(self)
    304                 self._fire_event(Events.ITERATION_STARTED)
    305                 self.state.output = self._process_function(self, batch)
--> 306                 self._fire_event(Events.ITERATION_COMPLETED)
    307                 if self.should_terminate or self.should_terminate_single_epoch:
    308                     self.should_terminate_single_epoch = False

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/engine/engine.py in _fire_event(self, event_name, *event_args, **event_kwargs)
    257             for func, args, kwargs in self._event_handlers[event_name]:
    258                 kwargs.update(event_kwargs)
--> 259                 func(self, *(event_args + args), **kwargs)
    260 
    261     def fire_event(self, event_name):

/anaconda3/envs/ignite/lib/python3.7/site-packages/pytorch_ignite-0.2.0-py3.7.egg/ignite/contrib/handlers/tensorboard_logger.py in __call__(self, engine, logger, event_name)
    219             name = name.replace('.', '/')
    220             logger.writer.add_scalar("grads_{}/{}".format(self.reduction.__name__, name),
--> 221                                      self.reduction(p.grad),
    222                                      global_step)
    223 

/anaconda3/envs/ignite/lib/python3.7/site-packages/torch/functional.py in norm(input, p, dim, keepdim, out, dtype)
    676         (tensor(3.7417), tensor(11.2250))
    677     """
--> 678     ndim = input.dim()
    679 
    680     # catch default case

AttributeError: 'NoneType' object has no attribute 'dim'
```


Discussion:

Patch:
diff --git a/ignite/contrib/handlers/tensorboard_logger.py b/ignite/contrib/handlers/tensorboard_logger.py
index 39734a42..92afdaee 100644
--- a/ignite/contrib/handlers/tensorboard_logger.py
+++ b/ignite/contrib/handlers/tensorboard_logger.py
@@ -140,6 +140,9 @@ class WeightsScalarHandler(BaseWeightsScalarHandler):
 
         global_step = engine.state.get_event_attrib_value(event_name)
         for name, p in self.model.named_parameters():
+            if p.grad is None:
+                continue
+
             name = name.replace('.', '/')
             logger.writer.add_scalar("weights_{}/{}".format(self.reduction.__name__, name),
                                      self.reduction(p.data),
@@ -177,6 +180,9 @@ class WeightsHistHandler(BaseWeightsHistHandler):
 
         global_step = engine.state.get_event_attrib_value(event_name)
         for name, p in self.model.named_parameters():
+            if p.grad is None:
+                continue
+
             name = name.replace('.', '/')
             logger.writer.add_histogram(tag="weights/{}".format(name),
                                         values=p.data.detach().cpu().numpy(),
@@ -216,6 +222,9 @@ class GradsScalarHandler(BaseWeightsScalarHandler):
 
         global_step = engine.state.get_event_attrib_value(event_name)
         for name, p in self.model.named_parameters():
+            if p.grad is None:
+                continue
+
             name = name.replace('.', '/')
             logger.writer.add_scalar("grads_{}/{}".format(self.reduction.__name__, name),
                                      self.reduction(p.grad),
@@ -252,6 +261,9 @@ class GradsHistHandler(BaseWeightsHistHandler):
 
         global_step = engine.state.get_event_attrib_value(event_name)
         for name, p in self.model.named_parameters():
+            if p.grad is None:
+                continue
+
             name = name.replace('.', '/')
             logger.writer.add_histogram(tag="grads/{}".format(name),
                                         values=p.grad.detach().cpu().numpy(),


Test patch:
diff --git a/tests/ignite/contrib/handlers/conftest.py b/tests/ignite/contrib/handlers/conftest.py
new file mode 100644
index 00000000..8f51d777
--- /dev/null
+++ b/tests/ignite/contrib/handlers/conftest.py
@@ -0,0 +1,44 @@
+import numpy as np
+import pytest
+import torch
+from mock import Mock
+
+
+@pytest.fixture()
+def norm_mock():
+    def norm(x):
+        return np.linalg.norm(x)
+    norm_mock = Mock(side_effect=norm, spec=norm)
+    norm_mock.configure_mock(**{"__name__": "norm"})
+    norm_mock.reset_mock()
+    return norm_mock
+
+
+@pytest.fixture()
+def dummy_model_factory():
+    class DummyModel(torch.nn.Module):
+        def __init__(self):
+            super(DummyModel, self).__init__()
+            self.fc1 = torch.nn.Linear(10, 10)
+            self.fc2 = torch.nn.Linear(12, 12)
+            self.fc1.weight.data.zero_()
+            self.fc1.bias.data.zero_()
+            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):
+        model = DummyModel()
+        if with_grads:
+            model.fc2.weight.grad = torch.zeros_like(model.fc2.weight)
+            model.fc2.bias.grad = torch.zeros_like(model.fc2.bias)
+
+            if not with_frozen_layer:
+                model.fc1.weight.grad = torch.zeros_like(model.fc1.weight)
+                model.fc1.bias.grad = torch.zeros_like(model.fc1.bias)
+
+        if with_frozen_layer:
+            for param in model.fc1.parameters():
+                param.requires_grad = False
+        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
index 5ceb0cfa..f32d3843 100644
--- a/tests/ignite/contrib/handlers/test_tensorboard_logger.py
+++ b/tests/ignite/contrib/handlers/test_tensorboard_logger.py
@@ -3,6 +3,7 @@ import tempfile
 import shutil
 import math
 
+import numpy as np
 import pytest
 
 from mock import MagicMock, call, ANY
@@ -183,20 +184,9 @@ def test_weights_scalar_handler_wrong_setup():
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
-def test_weights_scalar_handler():
+def test_weights_scalar_handler(dummy_model_factory):
 
-    class DummyModel(torch.nn.Module):
-
-        def __init__(self):
-            super(DummyModel, self).__init__()
-            self.fc1 = torch.nn.Linear(10, 10)
-            self.fc2 = torch.nn.Linear(12, 12)
-            self.fc1.weight.data.zero_()
-            self.fc1.bias.data.zero_()
-            self.fc2.weight.data.fill_(1.0)
-            self.fc2.bias.data.fill_(1.0)
-
-    model = DummyModel()
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
     wrapper = WeightsScalarHandler(model)
     mock_logger = MagicMock(spec=TensorboardLogger)
@@ -217,6 +207,34 @@ def test_weights_scalar_handler():
     ], any_order=True)
 
 
+def test_weights_scalar_handler_frozen_layers(dummy_model_factory):
+
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=True)
+
+    wrapper = WeightsScalarHandler(model)
+    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_scalar.assert_has_calls([
+        call("weights_norm/fc2/weight", 12.0, 5),
+        call("weights_norm/fc2/bias", math.sqrt(12.0), 5),
+    ], any_order=True)
+
+    with pytest.raises(AssertionError):
+        mock_logger.writer.add_scalar.assert_has_calls([
+            call("weights_norm/fc1/weight", 12.0, 5),
+            call("weights_norm/fc1/bias", math.sqrt(12.0), 5),
+        ], any_order=True)
+
+    assert mock_logger.writer.add_scalar.call_count == 2
+
+
 def test_weights_hist_handler_wrong_setup():
 
     with pytest.raises(TypeError, match="Argument model should be of type torch.nn.Module"):
@@ -230,20 +248,9 @@ def test_weights_hist_handler_wrong_setup():
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
-def test_weights_hist_handler():
-
-    class DummyModel(torch.nn.Module):
-
-        def __init__(self):
-            super(DummyModel, self).__init__()
-            self.fc1 = torch.nn.Linear(10, 10)
-            self.fc2 = torch.nn.Linear(12, 12)
-            self.fc1.weight.data.zero_()
-            self.fc1.bias.data.zero_()
-            self.fc2.weight.data.fill_(1.0)
-            self.fc2.bias.data.fill_(1.0)
+def test_weights_hist_handler(dummy_model_factory):
 
-    model = DummyModel()
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
     wrapper = WeightsHistHandler(model)
     mock_logger = MagicMock(spec=TensorboardLogger)
@@ -264,6 +271,33 @@ def test_weights_hist_handler():
     ], any_order=True)
 
 
+def test_weights_hist_handler_frozen_layers(dummy_model_factory):
+
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=True)
+
+    wrapper = WeightsHistHandler(model)
+    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_has_calls([
+        call(tag="weights/fc2/weight", values=ANY, global_step=5),
+        call(tag="weights/fc2/bias", values=ANY, global_step=5),
+    ], any_order=True)
+
+    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)
+    assert mock_logger.writer.add_histogram.call_count == 2
+
+
 def test_grads_scalar_handler_wrong_setup():
 
     with pytest.raises(TypeError, match="Argument model should be of type torch.nn.Module"):
@@ -280,43 +314,57 @@ def test_grads_scalar_handler_wrong_setup():
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
-def test_grads_scalar_handler():
+def test_grads_scalar_handler(dummy_model_factory, norm_mock):
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
-    class DummyModel(torch.nn.Module):
+    wrapper = GradsScalarHandler(model, reduction=norm_mock)
+    mock_logger = MagicMock(spec=TensorboardLogger)
+    mock_logger.writer = MagicMock()
+
+    mock_engine = MagicMock()
+    mock_engine.state = State()
+    mock_engine.state.epoch = 5
+    norm_mock.reset_mock()
 
-        def __init__(self):
-            super(DummyModel, self).__init__()
-            self.fc1 = torch.nn.Linear(10, 10)
-            self.fc2 = torch.nn.Linear(12, 12)
-            self.fc1.weight.data.zero_()
-            self.fc1.bias.data.zero_()
-            self.fc2.weight.data.fill_(1.0)
-            self.fc2.bias.data.fill_(1.0)
+    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+
+    mock_logger.writer.add_scalar.assert_has_calls([
+        call("grads_norm/fc1/weight", ANY, 5),
+        call("grads_norm/fc1/bias", ANY, 5),
+        call("grads_norm/fc2/weight", ANY, 5),
+        call("grads_norm/fc2/bias", ANY, 5),
+    ], any_order=True)
+    assert mock_logger.writer.add_scalar.call_count == 4
+    assert norm_mock.call_count == 4
 
-    model = DummyModel()
 
-    def norm(x):
-        return 0.0
+def test_grads_scalar_handler_frozen_layers(dummy_model_factory, norm_mock):
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=True)
 
-    wrapper = GradsScalarHandler(model, reduction=norm)
+    wrapper = GradsScalarHandler(model, reduction=norm_mock)
     mock_logger = MagicMock(spec=TensorboardLogger)
     mock_logger.writer = MagicMock()
 
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 5
+    norm_mock.reset_mock()
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    assert mock_logger.writer.add_scalar.call_count == 4
-
     mock_logger.writer.add_scalar.assert_has_calls([
-        call("grads_norm/fc1/weight", ANY, 5),
-        call("grads_norm/fc1/bias", ANY, 5),
         call("grads_norm/fc2/weight", ANY, 5),
         call("grads_norm/fc2/bias", ANY, 5),
     ], any_order=True)
 
+    with pytest.raises(AssertionError):
+        mock_logger.writer.add_scalar.assert_has_calls([
+            call("grads_norm/fc1/weight", ANY, 5),
+            call("grads_norm/fc1/bias", ANY, 5),
+        ], any_order=True)
+    assert mock_logger.writer.add_scalar.call_count == 2
+    assert norm_mock.call_count == 2
+
 
 def test_grads_hist_handler_wrong_setup():
 
@@ -331,20 +379,8 @@ def test_grads_hist_handler_wrong_setup():
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
-def test_grads_hist_handler():
-
-    class DummyModel(torch.nn.Module):
-
-        def __init__(self):
-            super(DummyModel, self).__init__()
-            self.fc1 = torch.nn.Linear(10, 10)
-            self.fc2 = torch.nn.Linear(12, 12)
-            self.fc1.weight.data.zero_()
-            self.fc1.bias.data.zero_()
-            self.fc2.weight.data.fill_(1.0)
-            self.fc2.bias.data.fill_(1.0)
-
-    model = DummyModel()
+def test_grads_hist_handler(dummy_model_factory):
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
     wrapper = GradsHistHandler(model)
     mock_logger = MagicMock(spec=TensorboardLogger)
@@ -354,12 +390,6 @@ def test_grads_hist_handler():
     mock_engine.state = State()
     mock_engine.state.epoch = 5
 
-    model.fc1.weight.grad = torch.zeros_like(model.fc1.weight)
-    model.fc1.bias.grad = torch.zeros_like(model.fc1.bias)
-
-    model.fc2.weight.grad = torch.zeros_like(model.fc2.weight)
-    model.fc2.bias.grad = torch.zeros_like(model.fc2.bias)
-
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
     assert mock_logger.writer.add_histogram.call_count == 4
@@ -371,6 +401,32 @@ def test_grads_hist_handler():
     ], any_order=True)
 
 
+def test_grads_hist_frozen_layers(dummy_model_factory):
+    model = dummy_model_factory(with_grads=True, with_frozen_layer=True)
+
+    wrapper = GradsHistHandler(model)
+    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)
+
+    assert mock_logger.writer.add_histogram.call_count == 2
+    mock_logger.writer.add_histogram.assert_has_calls([
+        call(tag="grads/fc2/weight", values=ANY, global_step=5),
+        call(tag="grads/fc2/bias", values=ANY, global_step=5),
+    ], any_order=True)
+
+    with pytest.raises(AssertionError):
+        mock_logger.writer.add_histogram.assert_has_calls([
+            call(tag="grads/fc1/weight", values=ANY, global_step=5),
+            call(tag="grads/fc1/bias", values=ANY, global_step=5),
+        ], any_order=True)
+
+
 def test_integration(dirname):
 
     n_epochs = 5