Package ai.djl.training.listener

Class SaveModelTrainingListener

java.lang.Object

ai.djl.training.listener.TrainingListenerAdapter

ai.djl.training.listener.SaveModelTrainingListener

All Implemented Interfaces:

TrainingListener

public class SaveModelTrainingListener
extends TrainingListenerAdapter

A TrainingListener that saves a model and can save checkpoints.

Nested Class Summary

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

TrainingListener.BatchData, TrainingListener.Defaults

Constructor Summary

Constructors

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

Method Summary

Modifier and Type	Method	Description
int	getCheckpoint()	Returns the checkpoint frequency (or -1 for no checkpointing) in SaveModelTrainingListener.
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	setCheckpoint(int checkpoint)	Sets the checkpoint frequency in SaveModelTrainingListener.
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

SaveModelTrainingListener

public SaveModelTrainingListener(java.lang.String outputDir)

Constructs a SaveModelTrainingListener using the model's name.

Parameters:

outputDir - the directory to output the checkpointed models in

SaveModelTrainingListener

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

Constructs a SaveModelTrainingListener.

Parameters:

overrideModelName - an override model name to save checkpoints with

outputDir - the directory to output the checkpointed models in

SaveModelTrainingListener

public SaveModelTrainingListener(java.lang.String outputDir,
                                 java.lang.String overrideModelName,
                                 int checkpoint)

Constructs a SaveModelTrainingListener.

Parameters:

overrideModelName - an override model name to save checkpoints with

outputDir - the directory to output the checkpointed models in

checkpoint - adds a checkpoint every n epochs

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

getCheckpoint

public int getCheckpoint()

Returns the checkpoint frequency (or -1 for no checkpointing) in SaveModelTrainingListener.

Returns:

the checkpoint frequency (or -1 for no checkpointing)

setCheckpoint

public void setCheckpoint(int checkpoint)

Sets the checkpoint frequency in SaveModelTrainingListener.

Parameters:

checkpoint - how many epochs between checkpoints (or -1 for no checkpoints)

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)