[Feat] init callbacks

[shijianjian]

We add the feature for Callbacks for augmentations.

As a start, we aim to support logging the transformed images for debugging purposes, on platforms like wandb, tensorboard, and local folders.

related to #941

[johnnv1]

can't we add these start/end in a decorator call because, in some cases, we'll need a "setup" before the hooks?

[shijianjian]

For the logging purpose, we may use a decorator. But if we want to log more things, this way can provide a better control.

There are some setups for different functions, like to compute the parameters when it is None, etc.

I am not quite sure what is the best practice for this here. If you got some better designs, we can do it in this PR.

[johnnv1]

Well, my first thought is to split these functions into something like

def forward(...):
    ... = self._setup_forward() # compute parameter, convert anything, etc
    self._forward_start_hooks() # execute hooks before starting the method itself
    ... = self._forward() # execute the implementation itself
    self._forward_end_hooks() # execute hooks after finishing the method itself
    return ...

But this may overcomplicate the API, and make it pretty unclear / splited -- so I'm not sure about it

---

For calling the hooks itself, instead of adding these inline loops, we may have something like:

def _run_callbacks_hook(callbacks: List[AugmentationCallbackBase], hook: str, *args, **kwargs) -> None:
    for cb in callbacks:
        if not hasattr(cb, hook):
            continue

        hook_callable = getattr(cb, hook)

        if not callable(hook_callable):
            continue

        hook_callable(*args, **kwargs)

@@ -208,10 +218,10 @@ class _BasicAugmentationBase(Module):
 
         params, flags = self._process_kwargs_to_params_and_flags(params, self.flags, **kwargs)
 
-        [cb.on_forward_start(input, params, flags) for cb in self.callbacks]
+        _run_callbacks_hook(self.callbacks, "on_forward_start", input, params, flags)
         output = self.apply_func(in_tensor, params, flags)
         output = self.transform_output_tensor(output, input_shape) if self.keepdim else output
-        [cb.on_forward_end(output, params, flags) for cb in self.callbacks]
+        _run_callbacks_hook(self.callbacks, "on_forward_end", output, params, flags)
         return output

[johnnv1]

btw, if we consider having _<method>_start_hooks, the setup itself can be considered a hook too

[shijianjian]

Ah. right. We can do it with this approach. It looks cleaner.

[edgarriba]

im a bit confused here since i see that this callback could be decoupled in two

1. AugmentationSequentialCallback (a good question is also if with callback we can change state and make nonlinear states)
   • handles the pipeline state like the forward
2. AugmentationOpCallback

• this specific to single augmentation.
• but e.g you could setup the callbacks so that you can inject things in between the pipeline

I think also that the callback should be able to forward propagate results and have access to the main workflow state

[shijianjian]

I was thinking the same, but it is somehow weird for users when call the API. Say, for logging images, a user needs to use:

AugmentationSequential(..., callbacks=[WandbSequentialLogger])

and

ColorJitter(..., callbacks=[WandbLogger])

Also, we are implementing something twice for very similar purposes.

I think also that the callback should be able to forward propagate results and have access to the main workflow state

I think we can pass the self object into the callbacks.

[johnnv1]

well, if I'm defining in the container, I would expect to have logs for the whole pipeline, and on the ops level, only for that op. But the public callback should be the same for both... otherwise, we'll overcomplicate the API usage

AugmentationSequential(
    Normalize(..., callbacks=[WandbLogger])
    ColorJitter(..., callbacks=None),
    RandomAffine(..., callbacks=[WandbLogger]),
    RandomSolarize(..., callbacks=None),
    ...,
    callbacks=[WandbLogger]
)

I would like it to work within ImageSequential which I think is a pretty nice feature, but we don't stand out as much

[edgarriba]

Sure, my point was more about that you might want a custom callback e.g to pipeline a variable from from normalize to random affine. The only dangerous thing I see now is that how to prevent that the callback in the Normalize implements something on the hooks belonging to the sequential commands. At internal level I think at least we should distinguish that

[shijianjian]

I am wondering if we actually need an augmentation-wise level of logging, or we support it only for the sequential wrapper.

[johnnv1]

yeap, i'm not sure when/why i would need that aside from debugging the aug