Use weak references to resolve circular references and add test

Author: goanpeca

ignite/engine/engine.py

@@ -328,7 +328,7 @@ def execute_something():
 
         try:
             _check_signature(handler, "handler", self, *(event_args + args), **kwargs)
-            self._event_handlers[event_name].append((handler, (self,) + args, kwargs))
+            self._event_handlers[event_name].append((handler, (weakref.ref(self),) + args, kwargs))
         except ValueError:
             _check_signature(handler, "handler", *(event_args + args), **kwargs)
             self._event_handlers[event_name].append((handler, args, kwargs))
@@ -432,7 +432,15 @@ def _fire_event(self, event_name: Any, *event_args: Any, **event_kwargs: Any) ->
         self.last_event_name = event_name
         for func, args, kwargs in self._event_handlers[event_name]:
             kwargs.update(event_kwargs)
-            first, others = ((args[0],), args[1:]) if (args and args[0] == self) else ((), args)
+            if args and isinstance(args[0], weakref.ref):
+                resolved_engine = args[0]()
+                if resolved_engine is None:
+                    raise RuntimeError("Engine reference not resolved. Cannot execute event handler.")
+                first, others = ((resolved_engine,), args[1:])
+            else:
+                # metrics do not provide engine when registered
+                first, others = (tuple(), args)  # type: ignore[assignment]
+
             func(*first, *(event_args + others), **kwargs)
 
     def fire_event(self, event_name: Any) -> None:

tests/ignite/engine/test_memory_leaks.py

@@ -0,0 +1,80 @@
+import sys
+import weakref
+
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+from torch.utils.data import DataLoader, TensorDataset
+
+from ignite.engine import create_supervised_trainer, create_supervised_evaluator, Events
+from ignite.handlers import ProgressBar, TensorboardLogger
+from ignite.handlers.tensorboard_logger import OptimizerParamsHandler
+from ignite.metrics import Loss
+
+
+class TestEngineMemoryLeak:
+    """See: https://github.com/pytorch/ignite/issues/3438"""
+
+    ENGINE_WEAK_REFS = {}
+
+    def do(self, model, dataloader, device, runs_folder):
+        optim = Adam(model.parameters(), 1e-4)
+        loss = nn.BCEWithLogitsLoss()
+        trainer = create_supervised_trainer(model, optim, loss, device)
+        metrics = {"Loss": Loss(loss)}
+        evaluator = create_supervised_evaluator(model, metrics, device)
+
+        pbar = ProgressBar()
+        pbar.attach(trainer)
+
+        tb_logger = TensorboardLogger(log_dir=runs_folder)
+        tb_logger.attach(trainer, OptimizerParamsHandler(optim), Events.EPOCH_STARTED)
+
+        trainer.run(dataloader, 1)
+
+        @trainer.on(Events.COMPLETED)
+        def completed(engine):
+            evaluator.run(dataloader)
+
+        tb_logger.close()
+        pbar.close()
+
+        self.ENGINE_WEAK_REFS[weakref.ref(trainer)] = sys.getrefcount(trainer) - 1
+        self.ENGINE_WEAK_REFS[weakref.ref(evaluator)] = sys.getrefcount(evaluator) - 1
+
+    def test_circular_references(self, tmp_path):
+        all_mem = []
+        all_max_mem = []
+        runs_folder = tmp_path / "runs"
+        runs_folder.mkdir()
+
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        else:
+            device = torch.device("cpu")
+
+        x = torch.rand(32, 1, 64, 64, 32)
+        y = torch.round(torch.rand(32, 1))
+        ds = TensorDataset(x, y)
+        dataloader = DataLoader(ds, 6)
+        for i in range(5):
+            N = 3000
+            model = nn.Sequential(nn.Flatten(), nn.Linear(64 * 64 * 32, N), nn.ReLU(), nn.Linear(N, 1))
+            model = model.to(device)
+            self.do(model, dataloader, device, runs_folder)
+            for engine_weak_ref, val in self.ENGINE_WEAK_REFS.items():
+                engine = engine_weak_ref()
+                if engine is not None:
+                    ref_count = sys.getrefcount(engine) - 1
+                    error_message = f"Engine Memory Leak: {engine} - Ref Count: {ref_count}"
+                    print(error_message)
+                    assert ref_count == 0
+
+            mem, max_mem = torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
+            print("!!!", i, mem, max_mem)
+            if all_mem and all_max_mem:
+                assert mem <= all_mem[-1], f"Memory Leak: {mem} > {all_mem[-1]}"
+                assert max_mem <= all_max_mem[-1], f"Max Memory Leak: {max_mem} > {all_max_mem[-1]}"
+
+            all_mem.append(mem)
+            all_max_mem.append(max_mem)