Logger utils

WandbConfig

基类：BaseModel

注意：name 控制实验名称由 NeMoLogger 处理，因此此处省略。 directory 也被省略，因为它由 NeMoLogger 设置。

参数

名称	类型	描述	默认值
entity		发布此运行的团队（默认：您的用户名或您的默认团队）	必需
project		此运行所属的项目名称。	必需
tags		与此运行关联的标签。	必需
group		给定组中所有运行共享的唯一字符串	必需
job_type		运行类型，当您将运行分组到更大的实验中时，这非常有用。	必需
offline		离线运行（数据稍后可以流式传输到 wandb 服务器）。	必需
id		设置版本，主要用于恢复之前的运行。	必需
anonymous		启用或显式禁用匿名日志记录。	必需

源代码位于 bionemo/llm/utils/logger_utils.py

class WandbConfig(BaseModel):
    """Note: `name` controls the exp name is handled by the NeMoLogger so it is ommitted here.
    `directory` is also omitted since it is set by the NeMoLogger.

    Args:
        entity: The team posting this run (default: your username or your default team)
        project: The name of the project to which this run will belong.
        tags: Tags associated with this run.
        group: A unique string shared by all runs in a given group
        job_type: Type of run, which is useful when you're grouping runs together into larger experiments.
        offline: Run offline (data can be streamed later to wandb servers).
        id: Sets the version, mainly used to resume a previous run.
        anonymous: Enables or explicitly disables anonymous logging.
    """  # noqa: D205

    entity: str | None  # The team posting this run (default: your username or your default team)
    project: str  # The name of the project to which this run will belong.
    # name: #Display name for the run. "This is handled by NeMoLogger"
    # save_dir: #Path where data is saved. "This is handled by NeMoLogger"
    tags: List[str] | None  # Tags associated with this run.
    group: str | None  # A unique string shared by all runs in a given group.
    job_type: str | None = (
        None  # Type of run, which is useful when you're grouping runs together into larger experiments.
    )
    offline: bool  # Run offline (data can be streamed later to wandb servers).
    id: str | None  # Sets the version, mainly used to resume a previous run.
    anonymous: bool  # Enables or explicitly disables anonymous logging.
    log_model: bool  # Save checkpoints in wandb dir to upload on W&B servers.

setup_nemo_lightning_logger(name='default-name', root_dir='./results', initialize_tensorboard_logger=False, wandb_config=None, ckpt_callback=None, **kwargs)

设置实验的 logger。

参数

名称	类型	描述	默认值
name	str	实验的名称。结果将放入 root_dir/name	'default-name'
root_dir	str | Path	根目录，用于创建 name 目录以保存运行结果。	'./results'
initialize_tensorboard_logger	bool	是否初始化 tensorboard logger。	False
wandb_config	Optional[WandbConfig]	wandb logger 的其余配置选项。	None
ckpt_callback	Optional[ModelCheckpoint]	要使用的检查点回调，必须是 pytorch lightning ModelCheckpoint 回调的子类。注意：底层 NeMoCheckpoint 构造函数中的类型注释不正确。	None
**kwargs	Dict[str, Any]	NeMoLogger 的 kwargs。	{}

返回

名称	类型	描述
NeMoLogger	NeMoLogger	NeMo logger 实例。

源代码位于 bionemo/llm/utils/logger_utils.py

def setup_nemo_lightning_logger(
    name: str = "default-name",
    root_dir: str | pathlib.Path = "./results",
    initialize_tensorboard_logger: bool = False,
    wandb_config: Optional[WandbConfig] = None,
    ckpt_callback: Optional[nemo_callbacks.ModelCheckpoint] = None,
    **kwargs: Dict[str, Any],
) -> NeMoLogger:
    """Setup the logger for the experiment.

    Arguments:
        name: The name of the experiment. Results go into `root_dir`/`name`
        root_dir: The root directory to create the `name` directory in for saving run results.
        initialize_tensorboard_logger: Whether to initialize the tensorboard logger.
        wandb_config: The remaining configuration options for the wandb logger.
        ckpt_callback: The checkpoint callback to use, must be a child of the pytorch lightning ModelCheckpoint callback.
            NOTE the type annotation in the underlying NeMoCheckpoint constructor is incorrect.
        **kwargs: The kwargs for the NeMoLogger.

    Returns:
        NeMoLogger: NeMo logger instance.
    """
    # The directory that the logger will save to
    save_dir = pathlib.Path(root_dir) / name
    if wandb_config is not None:
        wandb_logger = WandbLogger(save_dir=save_dir, name=name, **wandb_config.model_dump())
    else:
        wandb_logger = None
        logging.warning("WandB is currently turned off.")
    if initialize_tensorboard_logger:
        tb_logger = TensorBoardLogger(save_dir=save_dir, name=name)
    else:
        tb_logger = None
        logging.warning("User-set tensorboard is currently turned off. Internally one may still be set by NeMo2.")
    logger: NeMoLogger = NeMoLogger(
        name=name,
        log_dir=str(root_dir),
        tensorboard=tb_logger,
        wandb=wandb_logger,
        ckpt=ckpt_callback,
        use_datetime_version=False,
        version="dev",
        **kwargs,
    )
    # Needed so that the trainer can find an output directory for the profiler
    logger.save_dir = save_dir
    return logger