From e36a4c94d3340ffd3dda7a52f904ef9a99b47baa Mon Sep 17 00:00:00 2001 From: Lilferrit Date: Thu, 5 Sep 2024 14:38:16 -0700 Subject: [PATCH] write tensorboard to output directory --- casanovo/config.py | 2 +- casanovo/config.yaml | 4 ++-- casanovo/denovo/model_runner.py | 15 +++++++++++++-- 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/casanovo/config.py b/casanovo/config.py index dae46606..e1670aa3 100644 --- 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 index 1fcc77b4..59df016b 100644 --- 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 index 73d1da77..9e94199e 100644 --- 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,