diff --git a/deeplabcut/pose_estimation_pytorch/config/logger.py b/deeplabcut/pose_estimation_pytorch/config/logger.py index a37a7d33c..de025590c 100644 --- a/deeplabcut/pose_estimation_pytorch/config/logger.py +++ b/deeplabcut/pose_estimation_pytorch/config/logger.py @@ -41,16 +41,27 @@ class WandbLoggerConfig(LoggerConfig): # type: Logger type (should be 'WandbLogger') project_name: The name of the wandb project run_name: The name of the wandb run + entity: The wandb user or team under which the run is logged + notes: A longer description of the run, displayed in the wandb UI + tags: Tags for the run, used to organize and filter runs in the wandb UI + group: The name of the group to which this run belongs + job_type: The type of job being logged (e.g. 'train' or 'eval') image_log_interval: How often train/test images are logged in epochs (if None, train/test inputs are never logged) model: The model architecture to log train_folder: The path of the folder containing training files. - wandb_kwargs: Additional keyword arguments to pass to wandb.init + wandb_kwargs: Additional keyword arguments to pass to wandb.init. Use this for + wandb.init parameters which are not declared above. """ type: Literal[LoggerType.WandbLogger] project_name: str = "deeplabcut" run_name: str = "tmp" + entity: str | None = None + notes: str | None = None + tags: list[str] | None = None + group: str | None = None + job_type: str | None = None image_log_interval: int | None = None model: dict | None = None train_folder: str | None = None diff --git a/deeplabcut/pose_estimation_pytorch/runners/logger.py b/deeplabcut/pose_estimation_pytorch/runners/logger.py index 22c747bf1..ed179aa39 100644 --- a/deeplabcut/pose_estimation_pytorch/runners/logger.py +++ b/deeplabcut/pose_estimation_pytorch/runners/logger.py @@ -275,7 +275,9 @@ def __init__( None, train/test inputs are never logged). model: The model to log. Defaults to None. train_folder: path to the train folder (used to store the W&B run identifiers) - wandb_kwargs: extra arguments to pass to ``wb.init`` + wandb_kwargs: extra arguments to pass to ``wb.init``. These can be given as + keyword arguments, or collected in a single ``wandb_kwargs`` mapping (as + declared by ``WandbLoggerConfig``). Example: logger = WandbLogger(project_name="mice", run_name="exp1", model=my_model) @@ -291,6 +293,9 @@ def __init__( if wandb.run is not None: wandb.finish() + # A nested ``wandb_kwargs``mapping must be flattened + wandb_kwargs.update(wandb_kwargs.pop("wandb_kwargs", None) or {}) + self.run = wandb.init( project=project_name, name=run_name,