Skip to content

Add auto logging of trainer metadata and optimizer state saving to help with resumation of training - #39

Open
wli51 wants to merge 3 commits into
WayScience:mainfrom
wli51:resume-training
Open

Add auto logging of trainer metadata and optimizer state saving to help with resumation of training#39
wli51 wants to merge 3 commits into
WayScience:mainfrom
wli51:resume-training

Conversation

@wli51

@wli51 wli51 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Also doing small refactors of modules modified by the feature addition to move code away from main module and move class methods as independent methods to facilitate easier unit testing.

wli51 added 3 commits August 6, 2026 11:56
…ate existing properties by best effort inferring some. Remove some confusing and unused properties. Adjust model saving defaults.
… trainer metadata and couple logging of model weights by logging of optimizer state to facilitate resuming of complete trainings in whole epoch increments.
save_best_model: bool = True,
) -> List[Path]:

if file_name_prefix is None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why change this file? isn't the old default equivalent behavior?

Comment on lines +17 to +18
if file_name_prefix is None:
file_name_prefix = 'optimizer'

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comment here, why not default to "optimizer" instead of None?

if trainer is None:
return

config = {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

worth adding epoch here?

mlflow.log_artifact(str(file_path), artifact_path=artifact_path)


def _log_trainer_artifact(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't catch this until reading the next file, but, to me, naming this file "_log_trainer_artifact" implies you're writing to a log file, but in the next file (and actually the function within) is saving the best model and logging artifacts.

Consider enhancing the name of this function

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