Tensorboard io (#374)

* file io console options * output console io options * file io options tests * changelog entry * revised changelog * file io console options * output console io options * file io options tests * changelog entry * revised changelog * Generate new screengrabs with rich-codex * requested changes * updated integration test * requested changes * Generate new screengrabs with rich-codex * requested changes, output setup refactor * ModelRunner documentation * requested changes, _setup_output unit test * write tensorboard to output directory * ModelRunner output root bug fix, setup_model documentation, sequence output setup bug fix * Generate new screengrabs with rich-codex * logging format character * write tensorboard to output directory * Generate new screengrabs with rich-codex --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Noble-Lab · Sep 16, 2024 · aefa73c · aefa73c
1 parent 4cff0c8
commit aefa73c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/casanovo/config.py b/casanovo/config.py
@@ -64,7 +64,7 @@ class Config:
         max_length=int,
         residues=dict,
         n_log=int,
-        tb_summarywriter=str,
+        tb_summarywriter=bool,
         train_label_smoothing=float,
         warmup_iters=int,
         cosine_schedule_period_iters=int,

diff --git a/casanovo/config.yaml b/casanovo/config.yaml
@@ -40,8 +40,8 @@ random_seed: 454
 # OUTPUT OPTIONS
 # Logging frequency in training steps.
 n_log: 1
-# Tensorboard directory to use for keeping track of training metrics.
-tb_summarywriter:
+# Whether to create tensorboard directory
+tb_summarywriter: false
 # Model validation and checkpointing frequency in training steps.
 val_check_interval: 50_000
 

diff --git a/casanovo/denovo/model_runner.py b/casanovo/denovo/model_runner.py
@@ -62,6 +62,7 @@ def __init__(
         """Initialize a ModelRunner"""
         self.config = config
         self.model_filename = model_filename
+        self.output_dir = output_dir
 
         # Initialized later:
         self.tmp_dir = None
@@ -268,6 +269,16 @@ def initialize_model(self, train: bool) -> None:
             Determines whether to set the model up for model training or
             evaluation / inference.
         """
+        tb_summarywriter = None
+        if self.config.tb_summarywriter:
+            if self.output_dir is None:
+                logger.warning(
+                    "Can not create tensorboard because the output directory "
+                    "is not set in the model runner."
+                )
+            else:
+                tb_summarywriter = self.output_dir / "tensorboard"
+
         model_params = dict(
             dim_model=self.config.dim_model,
             n_head=self.config.n_head,
@@ -284,7 +295,7 @@ def initialize_model(self, train: bool) -> None:
             n_beams=self.config.n_beams,
             top_match=self.config.top_match,
             n_log=self.config.n_log,
-            tb_summarywriter=self.config.tb_summarywriter,
+            tb_summarywriter=tb_summarywriter,
             train_label_smoothing=self.config.train_label_smoothing,
             warmup_iters=self.config.warmup_iters,
             cosine_schedule_period_iters=self.config.cosine_schedule_period_iters,
@@ -303,7 +314,7 @@ def initialize_model(self, train: bool) -> None:
             min_peptide_len=self.config.min_peptide_len,
             top_match=self.config.top_match,
             n_log=self.config.n_log,
-            tb_summarywriter=self.config.tb_summarywriter,
+            tb_summarywriter=tb_summarywriter,
             train_label_smoothing=self.config.train_label_smoothing,
             warmup_iters=self.config.warmup_iters,
             cosine_schedule_period_iters=self.config.cosine_schedule_period_iters,