diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2023-09-03 01:13:37 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2023-09-03 01:13:37 +0200 |
commit | a65d3ec18a5541cec5297769f1027422975a62bc (patch) | |
tree | 08e5e22f76db2449d265476f5fb42c5ea64a2776 /training/callbacks | |
parent | e4d618443808f0931bbef0b9e10a2c2a215281a5 (diff) |
Update confs and callbacks
Diffstat (limited to 'training/callbacks')
-rw-r--r-- | training/callbacks/wandb.py | 34 |
1 files changed, 3 insertions, 31 deletions
diff --git a/training/callbacks/wandb.py b/training/callbacks/wandb.py index d9bb9b8..6adbebe 100644 --- a/training/callbacks/wandb.py +++ b/training/callbacks/wandb.py @@ -62,32 +62,7 @@ class UploadConfigAsArtifact(Callback): experiment.use_artifact(artifact) -class UploadCheckpointsAsArtifact(Callback): - """Upload checkpoint to wandb as an artifact, at the end of a run.""" - - def __init__( - self, ckpt_dir: str = "checkpoints/", upload_best_only: bool = False - ) -> None: - self.ckpt_dir = Path(ckpt_dir) - self.upload_best_only = upload_best_only - - @rank_zero_only - def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None: - """Uploads model checkpoint to W&B.""" - logger = get_wandb_logger(trainer) - experiment = logger.experiment - ckpts = wandb.Artifact("experiment-ckpts", type="checkpoints") - - if self.upload_best_only: - ckpts.add_file(trainer.checkpoint_callback.best_model_path) - else: - for ckpt in (self.ckpt_dir).rglob("*.ckpt"): - ckpts.add_file(ckpt) - - experiment.use_artifact(ckpts) - - -class ImageToCaptionLogger(Callback): +class ImageToCaption(Callback): """Logs the image and output caption.""" def __init__(self, num_samples: int = 8, on_train: bool = True) -> None: @@ -114,7 +89,7 @@ class ImageToCaptionLogger(Callback): pl_module: LightningModule, outputs: dict, batch: Tuple[Tensor, Tensor], - batch_idx: int, + *args, ) -> None: """Logs predictions on validation batch end.""" if self.has_metrics(outputs): @@ -127,9 +102,7 @@ class ImageToCaptionLogger(Callback): pl_module: LightningModule, outputs: dict, batch: Tuple[Tensor, Tensor], - batch_idx: int, *args, - # dataloader_idx: int, ) -> None: """Logs predictions on validation batch end.""" if self.has_metrics(outputs): @@ -142,8 +115,7 @@ class ImageToCaptionLogger(Callback): pl_module: LightningModule, outputs: dict, batch: Tuple[Tensor, Tensor], - batch_idx: int, - dataloader_idx: int, + *args, ) -> None: """Logs predictions on train batch end.""" if self.has_metrics(outputs): |