ai.djl.training.listener

Class CheckpointsTrainingListener

java.lang.Object

ai.djl.training.listener.CheckpointsTrainingListener

All Implemented Interfaces:

TrainingListener

public class CheckpointsTrainingListener
extends java.lang.Object
implements TrainingListener

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	onTrainingBatch(Trainer trainer, TrainingListener.BatchData batchData) Listens to the end of training one batch of data during training.
void	onTrainingBegin(Trainer trainer) Listens to the beginning of training.
void	onTrainingEnd(Trainer trainer) Listens to the end of training.
void	onValidationBatch(Trainer trainer, TrainingListener.BatchData batchData) Listens to the end of validating one batch of data during validation.
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 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

Parameters:

trainer - the trainer the listener is attached to

onTrainingBatch

public void onTrainingBatch(Trainer trainer,
                            TrainingListener.BatchData batchData)

Listens to the end of training one batch of data during training.

Specified by:

onTrainingBatch in interface TrainingListener

Parameters:

trainer - the trainer the listener is attached to

batchData - the data from the batch

onValidationBatch

public void onValidationBatch(Trainer trainer,
                              TrainingListener.BatchData batchData)

Listens to the end of validating one batch of data during validation.

Specified by:

onValidationBatch in interface TrainingListener

Parameters:

trainer - the trainer the listener is attached to

batchData - the data from the batch

onTrainingBegin

public void onTrainingBegin(Trainer trainer)

Listens to the beginning of training.

Specified by:

onTrainingBegin in interface TrainingListener

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

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)