diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-07-05 00:22:27 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-07-05 00:22:27 +0200 |
commit | 80e9bed0dd9840ac0cc9de1c6c1be3b6fed90cf9 (patch) | |
tree | 1e8dec8051220b2673cfa254991165cc6820b203 /text_recognizer/callbacks/wandb_callbacks.py | |
parent | 437ba4e22b2dad2ca309085a2e97f33dd03eb642 (diff) |
Add new updates to model and networks, reimplementing with attr
Diffstat (limited to 'text_recognizer/callbacks/wandb_callbacks.py')
-rw-r--r-- | text_recognizer/callbacks/wandb_callbacks.py | 9 |
1 files changed, 9 insertions, 0 deletions
diff --git a/text_recognizer/callbacks/wandb_callbacks.py b/text_recognizer/callbacks/wandb_callbacks.py index 3936aaf..900c3b1 100644 --- a/text_recognizer/callbacks/wandb_callbacks.py +++ b/text_recognizer/callbacks/wandb_callbacks.py @@ -29,6 +29,9 @@ class WatchModel(Callback): log: str = attr.ib(default="gradients") log_freq: int = attr.ib(default=100) + def __attrs_pre_init__(self): + super().__init__() + def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: """Watches model weights with wandb.""" logger = get_wandb_logger(trainer) @@ -41,6 +44,9 @@ class UploadCodeAsArtifact(Callback): project_dir: Path = attr.ib(converter=Path) + def __attrs_pre_init__(self): + super().__init__() + def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None: """Uploads project code as an artifact.""" logger = get_wandb_logger(trainer) @@ -59,6 +65,9 @@ class UploadCheckpointAsArtifact(Callback): ckpt_dir: Path = attr.ib(converter=Path) upload_best_only: bool = attr.ib() + def __attrs_pre_init__(self): + super().__init__() + def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None: """Uploads model checkpoint to W&B.""" logger = get_wandb_logger(trainer) |