Lightning-AI / Lightning-AI/pytorch-lightning

`LightningModule.on_train_batch_end` executes after `Callback.on_train_batch_end`

Open
#18,265 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

callback hooks question refactor
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.