Skip to content

Logging config in wandb #166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wojsza05 merged 3 commits into from
Mar 4, 2026
Merged

Logging config in wandb #166

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
83d9cd8
Logging config in wandb
wojsza05 Feb 23, 2026
b995ec2
Formatting
wojsza05 Feb 23, 2026
f70aca2
Fix
wojsza05 Feb 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
load_training_state,
get_full_checkpoint_path,
)
from src.core.metric_loggers import NeptuneLogger, get_metric_logger
from src.core.metric_loggers import NeptuneLogger, WandbLogger, get_metric_logger
from src.core.model import Residual
import platform

Expand Down Expand Up @@ -216,6 +216,7 @@ def initialize_training_components(cfg: OmegaConf, metric_logger=None):
metric_logger = get_metric_logger(
metric_logger_config=cfg.infrastructure.metric_logger,
tracker_run_id=training_state["run_id"],
full_config=cfg,
)

# Other loggers do not have `run` method
Expand All @@ -238,6 +239,22 @@ def initialize_training_components(cfg: OmegaConf, metric_logger=None):
metric_logger.run["learning_rate"] = learning_rate
metric_logger.run["exp_lr"] = exp_lr

elif isinstance(metric_logger, WandbLogger) and (
training_state["run_id"] is None
or cfg.infrastructure.metric_logger.new_wandb_job
):
# Update wandb config
if metric_logger.run is not None:
metric_logger.run.log(
{
"learning_rate": learning_rate,
"exp_lr": exp_lr,
"full_save_checkpoints_path": get_full_checkpoint_path(
cfg.trainer.checkpoint.save.path
),
}
)

torch.manual_seed(cfg.trainer.train_dataloader.dataset.seed)

device = get_device()
Expand Down
7 changes: 7 additions & 0 deletions src/core/metric_loggers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@ def get_composition_file_path(hydra_config):
def get_metric_logger(
metric_logger_config: OmegaConf = None,
tracker_run_id: Optional[str] = None,
full_config: Optional[OmegaConf] = None,
):
_metric_logger = None
if metric_logger_config.type == "neptune":
Expand Down Expand Up @@ -187,13 +188,19 @@ def get_metric_logger(
wandb_run_id = None if metric_logger_config.new_wandb_job else tracker_run_id
rank = int(os.environ.get("RANK", 0))
if rank == 0:
wandb_config_dict = (
OmegaConf.to_container(full_config, resolve=True)
if full_config
else None
)
wandb_logger = wandb.init(
entity=metric_logger_config.wandb_entity,
project=metric_logger_config.project_name,
name=metric_logger_config.name,
tags=list(metric_logger_config.tags),
id=wandb_run_id,
resume="allow",
config=wandb_config_dict,
)
_metric_logger = WandbLogger(
run=wandb_logger, should_log=True, config=metric_logger_config
Expand Down