ai.djl.training.listener

Class CheckpointsTrainingListener

java.lang.Object

ai.djl.training.listener.TrainingListenerAdapter

ai.djl.training.listener.CheckpointsTrainingListener

All Implemented Interfaces:

TrainingListener

public class CheckpointsTrainingListener
extends TrainingListenerAdapter

A TrainingListener that saves a model checkpoint after each epoch.

Nested Class Summary

Nested classes/interfaces inherited from interface ai.djl.training.listener.TrainingListener

TrainingListener.BatchData, TrainingListener.Defaults

Constructor Summary

Constructors

Constructor and Description
CheckpointsTrainingListener(java.lang.String outputDir) Constructs a CheckpointsTrainingListener using the model's name.
CheckpointsTrainingListener(java.lang.String outputDir, java.lang.String overrideModelName) Constructs a CheckpointsTrainingListener.
CheckpointsTrainingListener(java.lang.String outputDir, java.lang.String overrideModelName, int step) Constructs a CheckpointsTrainingListener.

Method Summary

Modifier and Type	Method and Description
java.lang.String	getOverrideModelName() Returns the override model name to save checkpoints with.
void	onEpoch(Trainer trainer) Listens to the end of an epoch during training.
void	onTrainingEnd(Trainer trainer) Listens to the end of training.
protected void	saveModel(Trainer trainer)
void	setOverrideModelName(java.lang.String overrideModelName) Sets the override model name to save checkpoints with.
void	setSaveModelCallback(java.util.function.Consumer<Trainer> onSaveModel) Sets the callback function on model saving.

Methods inherited from class ai.djl.training.listener.TrainingListenerAdapter

onTrainingBatch, onTrainingBegin, onValidationBatch

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Detail

CheckpointsTrainingListener

public CheckpointsTrainingListener(java.lang.String outputDir)

Constructs a CheckpointsTrainingListener using the model's name.

Parameters:

outputDir - the directory to output the checkpointed models in

CheckpointsTrainingListener

public CheckpointsTrainingListener(java.lang.String outputDir,
                                   java.lang.String overrideModelName)

Constructs a CheckpointsTrainingListener.

Parameters:

overrideModelName - an override model name to save checkpoints with

outputDir - the directory to output the checkpointed models in

CheckpointsTrainingListener

public CheckpointsTrainingListener(java.lang.String outputDir,
                                   java.lang.String overrideModelName,
                                   int step)

Constructs a CheckpointsTrainingListener.

Parameters:

overrideModelName - an override model name to save checkpoints with

outputDir - the directory to output the checkpointed models in

step - the spacing between each checkpoint, use -1 to only save model at the end of training

Method Detail

onEpoch

public void onEpoch(Trainer trainer)

Listens to the end of an epoch during training.

Specified by:

onEpoch in interface TrainingListener

Overrides:

onEpoch in class TrainingListenerAdapter

Parameters:

trainer - the trainer the listener is attached to

onTrainingEnd

public void onTrainingEnd(Trainer trainer)

Listens to the end of training.

Specified by:

onTrainingEnd in interface TrainingListener

Overrides:

onTrainingEnd in class TrainingListenerAdapter

Parameters:

trainer - the trainer the listener is attached to

getOverrideModelName

public java.lang.String getOverrideModelName()

Returns the override model name to save checkpoints with.

Returns:

the override model name to save checkpoints with

setOverrideModelName

public void setOverrideModelName(java.lang.String overrideModelName)

Sets the override model name to save checkpoints with.

Parameters:

overrideModelName - the override model name to save checkpoints with

setSaveModelCallback

public void setSaveModelCallback(java.util.function.Consumer<Trainer> onSaveModel)

Sets the callback function on model saving.

This allows user to set custom properties to model metadata.

Parameters:

onSaveModel - the callback function on model saving

saveModel

protected void saveModel(Trainer trainer)