Skip to content

[FEATURE] Add epoch-level progress bar to full trainer#658

Open
eliasgomes wants to merge 1 commit into
sdatkinson:mainfrom
eliasgomes:feat/epoch-progress-bar
Open

[FEATURE] Add epoch-level progress bar to full trainer#658
eliasgomes wants to merge 1 commit into
sdatkinson:mainfrom
eliasgomes:feat/epoch-progress-bar

Conversation

@eliasgomes

@eliasgomes eliasgomes commented May 4, 2026

Copy link
Copy Markdown

When training with nam-full, only a per-batch progress bar is shown, making it hard to tell overall training progress across epochs. This adds a second Rich progress bar below the existing one that tracks epoch completion with elapsed and remaining time estimates.

Before

before

After

after

Changes

  • Add _EpochProgressBar callback extending RichProgressBar with a separate epoch-level progress bar
  • Add _EpochCountColumn for epoch count display using Lightning's N/max-1 convention
  • Use transient Rich Progress for clean terminal output on both normal completion and Ctrl+C interruption
  • Remove v_num from displayed metrics (always 0 for single-logger setups)
  • Show just Epoch N on the per-batch bar (total is redundant with the epoch bar)
  • Add unit tests for both new classes

@eliasgomes eliasgomes force-pushed the feat/epoch-progress-bar branch from def46da to 9fc1dfc Compare June 3, 2026 13:47
Add a second Rich progress bar showing overall epoch progress below the
existing per-batch bar. Shows elapsed time, remaining time estimate, and
epoch count using Lightning's N/max-1 convention.

- Extend RichProgressBar with _EpochProgressBar callback
- Add _EpochCountColumn for epoch count display
- Use transient Rich Progress for clean terminal output on both normal
  completion and Ctrl+C interruption
- Remove v_num from displayed metrics
- Add unit tests for both new classes
@eliasgomes eliasgomes force-pushed the feat/epoch-progress-bar branch from 9fc1dfc to 47c9776 Compare June 3, 2026 13:53
@sdatkinson

Copy link
Copy Markdown
Owner

This is a cool idea.

I'd like to support this--but optionally. I could've sworn I'd implemented the ability to add callbacks via the configuration (i.e. the learning.json).

So you could make my_package.my_module.MyCallback, and then in learning.json, add:

"callbacks": [
  {
    "name": "eliasgomes_package.nam.EpochProgressBar",
    "args": [],
    "kwargs": {}
  }
]

(probably under the "trainer" field as you've put "enable_progress_bar".)

And it'd initialize and append this (and maybe others) using nam.util.init.

If you want to PR me that then you can use your progress bar and I can put in other things I like 😉

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants