xlstm_jax.trainer.logger.file_logger

Attributes

LOGGER

Classes

FileLoggerConfig	Configuration for the file logger tool.
FileLogger

Module Contents

xlstm_jax.trainer.logger.file_logger.LOGGER

class xlstm_jax.trainer.logger.file_logger.FileLoggerConfig

Bases: xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig

Configuration for the file logger tool.

log_step_key

The key to use for the step in the logs.

log_epoch_key

The key to use for the epoch in the logs.

config_format

The format to use when logging the config.

log_dir

The directory to use for the logs. Is added to the log_path in the logger.

log_step_key: str = 'log_step'

log_epoch_key: str = 'log_epoch'

config_format: str = 'json'

log_dir: str = 'file_logs'

create(logger)

Creates the file logger tool.

Parameters:

logger (xlstm_jax.trainer.logger.base_logger.Logger)

Return type:

xlstm_jax.trainer.logger.base_logger.LoggerTool

class xlstm_jax.trainer.logger.file_logger.FileLogger(config, logger)

Bases: xlstm_jax.trainer.logger.base_logger.LoggerTool

Parameters:

• config (FileLoggerConfig)

• logger (xlstm_jax.trainer.logger.base_logger.Logger)

config

config_to_log = None

logger

log_path

logs

log_config(config)

Log the config to disk.

Parameters:

config (xlstm_jax.configs.ConfigDict | dict[str, xlstm_jax.configs.ConfigDict]) – The config to log.

setup()

Set up the file logger.

log_metrics(metrics, step, epoch, mode)

Log a single metric dictionary in the file logger.

The metrics are logged in a list and saved to disk at the end.

Parameters:

• metrics (xlstm_jax.common_types.HostMetrics) – The metrics to log.

• step (int) – The current step.

• epoch (int) – The current epoch.

• mode (str) – The mode of logging. Commonly “train”, “val”, or “test”.

finalize(status)

Finalize the file logger.

Writes out the logs to disk.

Parameters:

status (str) – The status of the training run (e.g. success, failure).