[Feature Request] "Removable" event handlers.

[asford]

I'm happy to open a PR for this immediately if the feature seems useful.

Ignite's event interface provides both add and remove operations however the "remove" operation requires a reference to the original event handler. This makes temporary or one-off event handlers somewhat tricky to implement.

Pytorch's existing hook interface provides a solution to this issue by returning a "RemovableHandle" handle from register_hook, holding a weakref to the target hook to support later removal.

Ignite's add_event_handler should implement a similar interface, returning a "RemovableEventHandle" rather than None. For example, in an attrs-style syntax:

@attr.s(auto_attribs=True)
class RemovableEventHandle:
    """Removable weak reference to an added event handler."""


    event_name: str
    handler: weakref.ref = attr.ib(converter=_as_weakref)
    engine: weakref.ref = attr.ib(converter=_as_weakref)

    def remove(self):
        """Remove handler from engine."""
        handler = self.handler()
        engine: Engine = self.engine()

        if handler and engine and engine.has_event_handler(handler):
            engine.remove_event_handler(handler, self.event_name)

    def __enter__(self):  # noqa: D
        return self

    def __exit__(self, type, value, tb):  # noqa: D
        self.remove()

Would support features like:

engine = Engine(process_function)

def print_epoch(engine):
    print("Epoch: {}".format(engine.state.epoch))

# Train with a temporary event handler.
with engine.add_event_handler(Events.EPOCH_COMPLETED, print_epoch):
    engine.run(dataloader, 10)

[vfdev-5]

@asford thanks for the feature request, interesting observation about missing event handle!
In my use-cases, such handle was event handler itself, such that I could remove the handler.

Providing a temporary event handler seems interesting idea, could you please provide more details for its use-cases ?

[asford]

I've implemented and used this feature to support single-iteration diagnostic handlers connected to an ignite engine. As a concrete example, I'm debugging model training via a handler that accumulates parameter gradients for a single iteration, logs gradient histograms to tensorboard after iteration completion, and then self-detaches from the engine via the removable handlers. This allows me to log gradient histograms on every n-th iteration in a relatively cleanly.

As a skeleton "single-shot" handler:

from contextlib import ContextStack

@attr.s
class LogGradientHandler:
    
    summary_writer = attr.ib()
    _handlers : ContextStack = attr.ib(factory=ContextStack, init=False)
    ... gradient monitoring state ...
    
    def for_iteration(self, engine, model):
        for name, param in dict(model.named_parameters()).items():
            parameter_hook = ...attach backward hook...
            self._handlers.enter_context(parameter_hook)
        
        self._handlers.enter_context(
            self.engine.add_event_handler(
                Events.ITERATION_COMPLETED, self._log_and_detach))

    def _log_and_detach(self, engine):
        ... summary_writer.add_histogram for stored_gradient in logged values ...
        # Remove backward hooks from parameters, unregister event.
        self._handlers.close()

model = ... create model ...
summary_writer = ... setup tensorboard ...
engine = ... setup engine ...

@engine.on(Events.ITERATION_STARTED)
def maybe_log_grads(self, engine):
    """ Log parameter gradient every 10th batch."""
    if engine.state.iteration % 10:
        LogGradientHandler(summary_writer).for_iteration(engine, model)

[vfdev-5]

So if I understand correctly this attach-remove behaviour is needed for executing a handler once n iterations ?
If it is the case, there is also another contrib possibility to perform this using custom events.
I wonder whether this could help in your situation ?

Another option how to use such self-detaching context manager could be with an evaluator on train and test datasets with specific handlers...

[asford]

Yes, there are multiple ways of implementing the equivalent behavior for single-shot handlers. I've just found this method to be low-overhead and relatively easy to keep tidy.

As you point out, I'm also using this feature to "reuse" an evaluator engine across multiple datasets with different logging configurations.

I think the context manger support offered by this feature is the key to increased usability.

[vfdev-5]

I think the context manger support offered by this feature is the key to increased usability.

@asford I agree with that ! If you wish to send a PR, it would be awesome :)
IMO, the class RemovableEventHandle from your example, looks good. Probably, we could integrate it without attrib as a new dependency... What do yout think ?

[asford]

Probably, we could integrate it without attrib as a new dependency... What do yout think ?

Absolutely, I'll open a PR with a vanilla implementation shortly.

[vfdev-5]

A side note, if we had some handler like LearningRateFinder it could be neatly used with this context manager:

with trainer.add_event_handler(Events.ITERATION_COMPLETED, LearningRateFinder()):
    trainer.run(train_loader, max_epoch=1)