Get the model from the best epoch?

[dpankova]

Hello! Is it correct that the model returned by the learner.fit is the model from the last epoch?
Is there a way to save the model from the epoch with the best loss?

[laggui]

Hello! Is it correct that the model returned by the learner.fit is the model from the last epoch?

That is correct! The model returned is from its latest state (i.e., the last epoch update).

Is there a way to save the model from the epoch with the best loss?

By default, the learner already keeps the best 2 checkpoints based on the validation loss

burn/crates/burn-train/src/learner/builder.rs

Lines 96 to 106 in 960a635

	checkpointer_strategy: Box::new(
	ComposedCheckpointingStrategy::builder()
	.add(KeepLastNCheckpoints::new(2))
	.add(MetricCheckpointingStrategy::new(
	&LossMetric::<B>::new(), // default to valid loss
	Aggregate::Mean,
	Direction::Lowest,
	Split::Valid,
	))
	.build(),
	),

You can load those from the {artifact_dir}/checkpoint folder.

Note that the checkpointing strategy can be customized with the builder by providing .with_checkpointing_strategy(strategy).

[dpankova]

Thank you!