| 48 | |
| 49 | |
| 50 | def init_wandb(config: _config.TrainConfig, *, resuming: bool, log_code: bool = False, enabled: bool = True): |
| 51 | if not enabled: |
| 52 | wandb.init(mode="disabled") |
| 53 | return |
| 54 | |
| 55 | ckpt_dir = config.checkpoint_dir |
| 56 | if not ckpt_dir.exists(): |
| 57 | raise FileNotFoundError(f"Checkpoint directory {ckpt_dir} does not exist.") |
| 58 | if resuming: |
| 59 | run_id = (ckpt_dir / "wandb_id.txt").read_text().strip() |
| 60 | wandb.init(id=run_id, resume="must", project=config.project_name) |
| 61 | else: |
| 62 | wandb.init( |
| 63 | name=config.exp_name, |
| 64 | config=dataclasses.asdict(config), |
| 65 | project=config.project_name, |
| 66 | ) |
| 67 | (ckpt_dir / "wandb_id.txt").write_text(wandb.run.id) |
| 68 | |
| 69 | if log_code: |
| 70 | wandb.run.log_code(epath.Path(__file__).parent.parent) |
| 71 | |
| 72 | |
| 73 | def _load_weights_and_validate(loader: _weight_loaders.WeightLoader, params_shape: at.Params) -> at.Params: |