xlstm_jax.trainer.logger.wandb_logger

xlstm_jax.trainer.logger.wandb_logger#

Attributes#

LOGGER

Classes#

WandBLoggerConfig

Configuration for the WandB logger tool.

WandBLogger

Module Contents#

xlstm_jax.trainer.logger.wandb_logger.LOGGER#
class xlstm_jax.trainer.logger.wandb_logger.WandBLoggerConfig#

Bases: xlstm_jax.trainer.logger.base_logger.LoggerToolsConfig

Configuration for the WandB logger tool.

wb_entity#

The WandB entity to log to.

wb_project#

The WandB project to log to.

wb_host#

The WandB host to log to.

wb_key#

The WandB API key to use. If None, the key will be read from the environment.

wb_name#

The name of the run.

wb_notes#

Notes to add to the run.

wb_settings#

Settings to pass to the WandB run.

wb_tags#

Tags to add to the run.

log_dir#

The directory to log to.

wb_entity: str = 'xlstm'#
wb_project: str = 'xlstm_nxai'#
wb_host: str = 'https://api.wandb.ai'#
wb_key: str | None = None#
wb_name: str | None = None#
wb_notes: str | None = None#
wb_settings: dict[str, Any]#
wb_tags: list[str] = []#
wb_resume_id: str | None = None#
log_dir: str = 'wandb'#
create(logger)#

Create a WandB 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.wandb_logger.WandBLogger(config, logger)#

Bases: xlstm_jax.trainer.logger.base_logger.LoggerTool

Parameters:
config#
config_to_log = None#
logger#
log_path#
wandb_run = None#
log_config(config)#

Log the config to WandB.

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

Parameters:

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

setup()#

Set up the WandB logger.

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

log_metrics(metrics, step, epoch, mode)#

Log a single metric dictionary in the WandB 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 current mode. Will be used as a prefix for the metrics.

finalize(status)#

Closes the WandB logger.

Parameters:

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