Lightning-AI / Lightning-AI/pytorch-lightning

Callback hook for on_after_optimizer_step

Open
#11,688 16 comments 13 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature hooks priority: 2
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🚀 Feature

A callback hook for `on_after_optimizer_step`.

### Motivation

There's a callback hook for `on_before_optimizer_step`, but not for `on_after_optimizer_step`.
That would be useful for implementing an [ExponentialMovingAverage](https://github.com/pytorch/vision/blob/b3cdec1fd808b287e4e0d31f989488c335df1812/references/classification/utils.py#L158) callback: I'd like to update the average weights [after](https://github.com/pytorch/vision/blob/b3cdec1fd808b287e4e0d31f989488c335df1812/references/classification/train.py#L52) the optimizer has updated the parameters. Doing this average weight update with `on_train_batch_end` hook will not be accurate, as the model weights may not get updated after every training batch (due to gradient accumulation).

cc @borda @tchaton @rohitgr7 @carmocca @awaelchli @ninginthecloud @daniellepintz

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 existing on_before_optimizer_step callback through the optimizer-step training loop. Compare it with the PyTorch Vision references/classification/utils.py and train.py links to understand the intended post-update timing. Done means the new hook runs after optimizer parameter updates, including with gradient accumulation, and supports the stated averaging use case.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.