utils.wandb_helper

 1import datetime
 2import logging
 3import os
 4import re
 5
 6import wandb
 7
 8from config.config import WANDB_ACTIVE
 9
10
11def wandb_init(
12    run_name,
13    project_name,
14    dataset_name,
15    config=None,
16    log_dir_root="logs/tensorboard/",
17    log_dir=None,
18    sync_tensorboard=True,
19    wandb_activate=True,
20):
21    """Initialize a Weights & Biases run for experiment tracking and logging with optional TensorBoard sync."""
22
23    # Logging dir
24    if log_dir is None:
25        project_folder = re.sub(r"[^\w\-_]", "_", project_name)
26        run_folder = re.sub(r"[^\w\-_]", "_", run_name)
27        now = datetime.datetime.now()
28        time_folder = now.strftime("%Y-%m-%d_%H-%M-%S")
29        log_dir = os.path.join(log_dir_root, project_folder, run_folder, time_folder)
30
31    wandb.tensorboard.patch(root_logdir=log_dir)
32
33    wandb_run = None
34    if wandb_activate and WANDB_ACTIVE:
35        wandb_run = wandb.init(
36            name=run_name,
37            sync_tensorboard=sync_tensorboard,
38            project=project_name,
39            tags=[dataset_name],
40            config=config,
41        )
42        logging.info(
43            f"Launched Weights and Biases run {wandb_run} with config {config}"
44        )
45    else:
46        logging.warning(
47            "WandB run was not initialized. Set 'wandb_activate' and WANDB_ACTIVE to True."
48        )
49
50    return wandb_run, log_dir
51
52
53def wandb_close(exit_code=0):
54    """Close active W&B run and unpatch tensorboard logging, returning exit code 0 by default."""
55
56    try:
57        logging.info(f"Closing Weights and Biases run with exit_code {exit_code}")
58        wandb.finish(exit_code=exit_code)
59    except:
60        logging.warning(f"WandB run could not be finished")
61    try:
62        wandb.tensorboard.unpatch()
63    except:
64        logging.warning(f"WandB tensorboard could not be unpatched")
def wandb_init( run_name, project_name, dataset_name, config=None, log_dir_root='logs/tensorboard/', log_dir=None, sync_tensorboard=True, wandb_activate=True):
12def wandb_init(
13    run_name,
14    project_name,
15    dataset_name,
16    config=None,
17    log_dir_root="logs/tensorboard/",
18    log_dir=None,
19    sync_tensorboard=True,
20    wandb_activate=True,
21):
22    """Initialize a Weights & Biases run for experiment tracking and logging with optional TensorBoard sync."""
23
24    # Logging dir
25    if log_dir is None:
26        project_folder = re.sub(r"[^\w\-_]", "_", project_name)
27        run_folder = re.sub(r"[^\w\-_]", "_", run_name)
28        now = datetime.datetime.now()
29        time_folder = now.strftime("%Y-%m-%d_%H-%M-%S")
30        log_dir = os.path.join(log_dir_root, project_folder, run_folder, time_folder)
31
32    wandb.tensorboard.patch(root_logdir=log_dir)
33
34    wandb_run = None
35    if wandb_activate and WANDB_ACTIVE:
36        wandb_run = wandb.init(
37            name=run_name,
38            sync_tensorboard=sync_tensorboard,
39            project=project_name,
40            tags=[dataset_name],
41            config=config,
42        )
43        logging.info(
44            f"Launched Weights and Biases run {wandb_run} with config {config}"
45        )
46    else:
47        logging.warning(
48            "WandB run was not initialized. Set 'wandb_activate' and WANDB_ACTIVE to True."
49        )
50
51    return wandb_run, log_dir

Initialize a Weights & Biases run for experiment tracking and logging with optional TensorBoard sync.

def wandb_close(exit_code=0):
54def wandb_close(exit_code=0):
55    """Close active W&B run and unpatch tensorboard logging, returning exit code 0 by default."""
56
57    try:
58        logging.info(f"Closing Weights and Biases run with exit_code {exit_code}")
59        wandb.finish(exit_code=exit_code)
60    except:
61        logging.warning(f"WandB run could not be finished")
62    try:
63        wandb.tensorboard.unpatch()
64    except:
65        logging.warning(f"WandB tensorboard could not be unpatched")

Close active W&B run and unpatch tensorboard logging, returning exit code 0 by default.