xlstm_jax.trainer.logger.tensorboard_logger

Attributes

LOGGER

Classes

TensorBoardLoggerConfig	Configuration for the TensorBoard logger tool.
TensorBoardLogger

Module Contents

xlstm_jax.trainer.logger.tensorboard_logger.LOGGER

class xlstm_jax.trainer.logger.tensorboard_logger.TensorBoardLoggerConfig

Bases: xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig

Configuration for the TensorBoard logger tool.

tb_flush_secs

The frequency at which to flush the tensorboard logs.

tb_max_queue

The maximum number of items to queue before flushing.

tb_new_style

Whether to use the new style of logging. See PyTorch SummaryWriter documentation for more information.

log_dir

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

tb_flush_secs: int = 120

tb_max_queue: int = 10

tb_new_style: bool = False

log_dir: str = 'tensorboard'

create(logger)

Creates the TensorBoard 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.tensorboard_logger.TensorBoardLogger(config, logger)

Bases: xlstm_jax.trainer.logger.base_logger.LoggerTool

Parameters:

• config (TensorBoardLoggerConfig)

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

config

config_to_log = None

logger

writer: torch.utils.tensorboard.SummaryWriter = None

log_config(config)

Log the config to TensorBoard.

If the writer is not set up, the config will be saved and logged when the writer is set up.

Parameters:

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

_log_config()

Logs stored config to TensorBoard if writer is set up.

setup()

Set up the TensorBoard logger.

If the writer is already set up, this function skips the setup.

log_metrics(metrics, step, epoch, mode)

Log a single metric dictionary in the TensorBoard logger.

Parameters:

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

• step (int) – The current step.

• epoch (int) – The current epoch. Currently unused.

• mode (str) – The mode of logging. Commonly “train”, “val”, or “test”. Is used as prefix for the metric keys.

finalize(status)

Close the TensorBoard logger.

Parameters:

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