Initialize wandb logging.
(config: _config.TrainConfig, *, resuming: bool, enabled: bool = True)
| 70 | |
| 71 | |
| 72 | def init_wandb(config: _config.TrainConfig, *, resuming: bool, enabled: bool = True): |
| 73 | """Initialize wandb logging.""" |
| 74 | if not enabled: |
| 75 | wandb.init(mode="disabled") |
| 76 | return |
| 77 | |
| 78 | ckpt_dir = config.checkpoint_dir |
| 79 | if not ckpt_dir.exists(): |
| 80 | raise FileNotFoundError(f"Checkpoint directory {ckpt_dir} does not exist.") |
| 81 | |
| 82 | if resuming: |
| 83 | run_id = (ckpt_dir / "wandb_id.txt").read_text().strip() |
| 84 | wandb.init(id=run_id, resume="must", project=config.project_name) |
| 85 | else: |
| 86 | wandb.init( |
| 87 | name=config.exp_name, |
| 88 | config=dataclasses.asdict(config), |
| 89 | project=config.project_name, |
| 90 | ) |
| 91 | (ckpt_dir / "wandb_id.txt").write_text(wandb.run.id) |
| 92 | |
| 93 | |
| 94 | def setup_ddp(): |