add log dir for keras train

calico · Apr 24, 2024 · 1cd10d7 · 1cd10d7
1 parent 9031b4e
commit 1cd10d7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/src/baskerville/scripts/hound_train.py b/src/baskerville/scripts/hound_train.py
@@ -56,6 +56,12 @@ def main():
         default="train_out",
         help="Output directory [Default: %(default)s]",
     )
+    parser.add_argument(
+        "-log",
+        "--log_dir",
+        default="log_out",
+        help="Log directory [Default: %(default)s]",
+    )
     parser.add_argument(
         "--restore",
         default=None,
@@ -150,7 +156,7 @@ def main():
 
         # initialize trainer
         seqnn_trainer = trainer.Trainer(
-            params_train, train_data, eval_data, args.out_dir
+            params_train, train_data, eval_data, args.out_dir, args.log_dir
         )
 
         # compile model
@@ -182,6 +188,7 @@ def main():
                 train_data,
                 eval_data,
                 args.out_dir,
+                args.log_dir,
                 strategy,
                 params_train["num_gpu"],
                 args.keras_fit,

diff --git a/src/baskerville/trainer.py b/src/baskerville/trainer.py
@@ -91,6 +91,7 @@ def __init__(
         train_data,
         eval_data,
         out_dir: str,
+        log_dir: str,
         strategy=None,
         num_gpu: int = 1,
         keras_fit: bool = False,
@@ -188,7 +189,7 @@ def fit_keras(self, seqnn_model):
 
         callbacks = [
             early_stop,
-            tf.keras.callbacks.TensorBoard(self.out_dir),
+            tf.keras.callbacks.TensorBoard(self.log_dir),
             tf.keras.callbacks.ModelCheckpoint("%s/model_check.h5" % self.out_dir),
             save_best,
         ]