Hi! In my
LightningModule, I have a
on_train_batch_end which performs some post-processing of the results from the
training_step (for example logits normalisation and update metrics). Now I can callbacks with also
on_train_batch_end method, which expects the model outputs to be normalised.
However, I find out that the execution order is “
LightningModule.on_train_batch_end”, while I was expecting to be “
Callback.on_train_batch_end”, which seems more natural to me.
I don’t want to have this normalisation process in the
training_step, because I want to have it as an
abstractumethod so children classes overwrite the
training_step and the
on_train_batch_end can be supported out-of-the-box.
Binary Segmentation model,
LightningModule.training_step produces the model
logits, in the
LightningModule.on_train_batch_end I define how these will be denormalized and pass them to the metrics, and I have a callback
SegmentationWriter which expects the predictions to be denomalized, to visualise and save them on disk. Note that I want to describe and execute the denormalisation process only once, so that the metrics and the segmentation writer will use the same one (and I dont want to repeat myself in the segmentation callback)
Question: is it intended or a specific reason which the
on_train_batch_end callback method come first from the
LightningModule, as for me it feels more natural the otherway around? Can I alter this order?