Package ai.djl.training.listener

Class MemoryTrainingListener

java.lang.Object

ai.djl.training.listener.TrainingListenerAdapter

ai.djl.training.listener.MemoryTrainingListener

All Implemented Interfaces:

TrainingListener

public class MemoryTrainingListener
extends TrainingListenerAdapter

TrainingListener that collects the memory usage information.

If an outputDir is provided, the file "$outputDir/memory.log" will be created after training with the memory usage results.

Nested Class Summary

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

TrainingListener.BatchData, TrainingListener.Defaults

Constructor Summary

Constructors

Constructor	Description
MemoryTrainingListener()	Constructs a MemoryTrainingListener that does not output data to a file.
MemoryTrainingListener(java.lang.String outputDir)	Constructs a MemoryTrainingListener that outputs data in the given directory.

Method Summary

Modifier and Type	Method	Description
static void	collectMemoryInfo(Metrics metrics)	Collects memory information.
static void	dumpMemoryInfo(Metrics metrics, java.lang.String logDir)	Dump memory metrics into log directory.
void	onTrainingBatch(Trainer trainer, TrainingListener.BatchData batchData)	Listens to the end of training one batch of data during 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.

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

onEpoch, onTrainingBegin

Methods inherited from class java.lang.Object

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

Constructor Detail

MemoryTrainingListener

public MemoryTrainingListener()

Constructs a MemoryTrainingListener that does not output data to a file.

MemoryTrainingListener

public MemoryTrainingListener(java.lang.String outputDir)

Constructs a MemoryTrainingListener that outputs data in the given directory.

If an output directory is provided, the file "$outputDir/memory.log" will be created after training with the memory usage results. The log file consists of heap bytes, non-heap bytes, cpu percentage and rss bytes consumption along with the timestamps.

Parameters:

outputDir - the directory to output the tracked memory data in

Method Detail

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

Overrides:

onTrainingBatch in class TrainingListenerAdapter

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

Overrides:

onValidationBatch in class TrainingListenerAdapter

Parameters:

trainer - the trainer the listener is attached to

batchData - the data from the batch

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

collectMemoryInfo

public static void collectMemoryInfo(Metrics metrics)

Collects memory information. In order to collect metrics, the Trainer must set metrics. Monitor the metrics by enabling the following flag in the command line arguments: -Dcollect-memory=true

Parameters:

metrics - Metrics to store memory information

dumpMemoryInfo

public static void dumpMemoryInfo(Metrics metrics,
                                  java.lang.String logDir)

Dump memory metrics into log directory.

Parameters:

metrics - metrics contains memory information

logDir - output log directory