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.

  • Constructor Details

    • MemoryTrainingListener

      public MemoryTrainingListener()
      Constructs a MemoryTrainingListener that does not output data to a file.

    • MemoryTrainingListener

      public MemoryTrainingListener(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 Details

    • 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, String logDir)
      Dump memory metrics into log directory.
      Parameters:
      metrics - metrics contains memory information
      logDir - output log directory