Lightning-AI / Lightning-AI/pytorch-lightning
`LightningModule.on_train_batch_end` executes after `Callback.on_train_batch_end`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Outline & Motivation
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.training_step` -> `Callback.on_train_batch_end` -> `LightningModule.on_train_batch_end`", while I was expecting to be "`LightningModule.training_step` -> `LightningModule.on_train_batch_end` -> `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.
Concreate example:
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?
Thank you!
### Pitch
_No response_
### Additional context
_No response_
cc @justusschock @awaelchli @carmocca @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by tracing the training-loop dispatch for LightningModule.on_train_batch_end and Callback.on_train_batch_end. Add a focused test that records both hook calls and verifies the intended order, then confirm the existing callback behavior and document what the test establishes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100