test

Author: goanpeca

ignite/engine/engine.py

@@ -17,6 +17,8 @@
 __all__ = ["Engine"]
 
 
+REF_COUNTER = 0
+
 class Engine(Serializable):
     """Runs a given ``process_function`` over each batch of a dataset, emitting events as it goes.
 
@@ -163,6 +165,21 @@ def __init__(self, process_function: Callable[["Engine", Any], Any]):
         # generator provided by self._internal_run_as_gen
         self._internal_run_generator: Optional[Generator[Any, None, State]] = None
 
+    def __del__(self) -> None:
+        """Finalize the engine."""
+        self._event_handlers.clear()
+        print("__del__", self)
+        print("Cleared event handlers for {}".format(self))
+    #     global REF_COUNTER
+    #     REF_COUNTER += 1
+    #     print("counter", REF_COUNTER)
+        # try to clear event handlers to break circular references
+        # try:
+        #     self._event_handlers.clear()
+        #     print("Cleared event handlers for {}".format(self))
+        # except Exception:
+        #     pass
+
     def register_events(
         self, *event_names: Union[List[str], List[EventEnum]], event_to_attr: Optional[dict] = None
     ) -> None:
@@ -328,6 +345,7 @@ def execute_something():
 
         try:
             _check_signature(handler, "handler", self, *(event_args + args), **kwargs)
+            # self._event_handlers[event_name].append((handler, (weakref.ref(self),) + args, kwargs))
             self._event_handlers[event_name].append((handler, (self,) + args, kwargs))
         except ValueError:
             _check_signature(handler, "handler", *(event_args + args), **kwargs)
@@ -433,6 +451,15 @@ def _fire_event(self, event_name: Any, *event_args: Any, **event_kwargs: Any) ->
         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_circular_references.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
+from ignite.handlers import ProgressBar, TensorboardLogger
+from ignite.handlers.tensorboard_logger import OptimizerParamsHandler
+from torch.optim import Adam
+from ignite.metrics import Loss
+from torch.utils.data import DataLoader, TensorDataset
+
+
+def do(model, dataloader, device):
+    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")
+    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()
+
+    # del trainer
+    # del evaluator
+
+def test_circular_references():
+    device = torch.device("cuda")
+    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)
+        do(model, dataloader, device)
+        print("!!!", i, torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated())