summaryrefslogtreecommitdiff
path: root/training/callbacks
diff options
context:
space:
mode:
Diffstat (limited to 'training/callbacks')
-rw-r--r--training/callbacks/wandb_callbacks.py12
1 files changed, 6 insertions, 6 deletions
diff --git a/training/callbacks/wandb_callbacks.py b/training/callbacks/wandb_callbacks.py
index 2264750..11d0936 100644
--- a/training/callbacks/wandb_callbacks.py
+++ b/training/callbacks/wandb_callbacks.py
@@ -36,19 +36,19 @@ class WatchModel(Callback):
logger.watch(model=trainer.model, log=self.log, log_freq=self.log_freq)
-class UploadCodeAsArtifact(Callback):
+class UploadConfigAsArtifact(Callback):
"""Upload all *.py files to W&B as an artifact, at the beginning of the run."""
def __init__(self) -> None:
- self.project_dir = Path(__file__).resolve().parents[2] / "text_recognizer"
+ self.config_dir = Path(".hydra/")
@rank_zero_only
def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
"""Uploads project code as an artifact."""
logger = get_wandb_logger(trainer)
experiment = logger.experiment
- artifact = wandb.Artifact("project-source", type="code")
- for filepath in self.project_dir.glob("**/*.py"):
+ artifact = wandb.Artifact("experiment-config", type="config")
+ for filepath in self.config_dir.rglob("*.yaml"):
artifact.add_file(str(filepath))
experiment.use_artifact(artifact)
@@ -60,7 +60,7 @@ class UploadCheckpointsAsArtifact(Callback):
def __init__(
self, ckpt_dir: str = "checkpoints/", upload_best_only: bool = False
) -> None:
- self.ckpt_dir = Path(__file__).resolve().parent / ckpt_dir
+ self.ckpt_dir = Path(ckpt_dir)
self.upload_best_only = upload_best_only
@rank_zero_only
@@ -73,7 +73,7 @@ class UploadCheckpointsAsArtifact(Callback):
if self.upload_best_only:
ckpts.add_file(trainer.checkpoint_callback.best_model_path)
else:
- for ckpt in (self.ckpt_dir).glob("**/*.ckpt"):
+ for ckpt in (self.ckpt_dir).rglob("**/*.ckpt"):
ckpts.add_file(ckpt)
experiment.use_artifact(ckpts)