Added debug mode for Engine

ignite/engine/engine.py

@@ -125,6 +125,10 @@ def compute_mean_std(engine, batch):
     _state_dict_all_req_keys = ("epoch_length", "max_epochs")
     _state_dict_one_of_opt_keys = ("iteration", "epoch")
 
+    DEBUG_EVENTS = 1
+    DEBUG_OUTPUT = 2
+    DEBUG_GRADS = 3
+
     # Flag to disable engine._internal_run as generator feature for BC
     interrupt_resume_enabled = True
 
@@ -143,6 +147,8 @@ def __init__(self, process_function: Callable[["Engine", Any], Any]):
         self._dataloader_iter: Optional[Iterator[Any]] = None
         self._init_iter: Optional[int] = None
 
+        self.debug_level = 0
+
         self.register_events(*Events)
 
         if self._process_function is None:
@@ -425,6 +431,27 @@ def _fire_event(self, event_name: Any, *event_args: Any, **event_kwargs: Any) ->
             first, others = ((args[0],), args[1:]) if (args and args[0] == self) else ((), args)
             func(*first, *(event_args + others), **kwargs)
 
+    def debug(self, level: int = 0, **kwargs: Any) -> None:
+        self.level = level
+
+        if self.level > 2:
+            self.lr = kwargs["optimizer"].param_groups[0]["lr"]
+            self.layer = kwargs["layer"]
+            self.logger.debug(
+                f"{self.state.epoch} | {self.state.iteration}, Firing handlers for event {self.last_event_name}, \
+                    Loss : {self.state.output}, LR : {self.lr}, Gradients : {self.layer.weight.grad}"
+            )
+        elif self.level > 1:
+            self.lr = kwargs["optimizer"].param_groups[0]["lr"]
+            self.logger.debug(
+                f"{self.state.epoch} | {self.state.iteration} Firing handlers for event {self.last_event_name}, \
+                    Loss : {self.state.output}, LR : {self.lr}"
+            )
+        elif self.level > 0:
+            self.logger.debug(
+                f"{self.state.epoch} | {self.state.iteration}, Firing handlers for event {self.last_event_name}"
+            )
+
     def fire_event(self, event_name: Any) -> None:
         """Execute all the handlers associated with given event.
 

tests/ignite/engine/test_engine.py

@@ -1388,3 +1388,72 @@ def check_iter_epoch():
     state = engine.run(data, max_epochs=max_epochs)
     assert state.iteration == max_epochs * len(data) and state.epoch == max_epochs
     assert num_calls_check_iter_epoch == 1
+
+
+def test_engine_debug():
+    import torch.nn.functional as F
+    from torch import nn
+    from torch.optim import SGD
+    from torch.utils.data import DataLoader
+    from torchvision.datasets import MNIST
+    from torchvision.transforms import Compose, ToTensor
+
+    from ignite.engine import create_supervised_trainer
+
+    DEBUG_EVENTS = 1
+    DEBUG_OUTPUT = 2
+    DEBUG_GRADS = 3
+
+    class Net(nn.Module):
+        def __init__(self):
+            super(Net, self).__init__()
+            self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+            self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+            self.conv2_drop = nn.Dropout2d()
+            self.fc1 = nn.Linear(320, 50)
+            self.fc2 = nn.Linear(50, 10)
+
+        def forward(self, x):
+            x = F.relu(F.max_pool2d(self.conv1(x), 2))
+            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+            x = x.view(-1, 320)
+            x = F.relu(self.fc1(x))
+            x = F.dropout(x, training=self.training)
+            x = self.fc2(x)
+            return F.log_softmax(x, dim=-1)
+
+    def _test():
+        train_loader = DataLoader(
+            MNIST(download=True, root=".", transform=Compose([ToTensor()]), train=True),
+            batch_size=64,
+            shuffle=True,
+        )
+
+        model = Net()
+        device = "cpu"
+        log_interval = 10
+        epochs = 10
+
+        if torch.cuda.is_available():
+            device = "cuda"
+
+        model.to(device)  # Move model before creating optimizer
+        optimizer = SGD(model.parameters(), lr=0.01, momentum=0.5)
+        criterion = nn.NLLLoss()
+        trainer = create_supervised_trainer(model, optimizer, criterion, device=device)
+
+        @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
+        def log_training_debug_events(engine):
+            trainer.debug(level=DEBUG_EVENTS)
+
+        @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
+        def log_training_debug_outputs(engine):
+            trainer.debug(level=DEBUG_OUTPUT, optimizer=optimizer)
+
+        @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
+        def log_training_debug_grads(engine):
+            trainer.debug(level=DEBUG_GRADS, optimizer=optimizer, layer=model.fc2)
+
+        trainer.run(train_loader, max_epochs=epochs)
+
+    _test()