modify train for composer

calico · Mar 6, 2024 · 121979b · 121979b
1 parent 2e58034
commit 121979b
Show file tree

Hide file tree

Showing 2 changed files with 574 additions and 0 deletions.
diff --git a/src/baskerville/scripts/hound_train.py b/src/baskerville/scripts/hound_train.py
@@ -16,6 +16,7 @@
 import argparse
 import json
 import os
+import sys
 import shutil
 
 import numpy as np
@@ -56,6 +57,12 @@ def main():
         default="train_out",
         help="Output directory [Default: %(default)s]",
     )
+    parser.add_argument(
+        "--outlog",
+        action="store_true",
+        default=False,
+        help="Keep a log file in out_dir/trainlog.txt [Default: %(default)s]",
+    )
     parser.add_argument(
         "--restore",
         default=None,
@@ -84,6 +91,9 @@ def main():
     )
     args = parser.parse_args()
 
+    if args.outlog:
+        orig_stdout = sys.stdout
+        sys.stdout = open("%s/trainlog.txt" % args.out_dir, "w")
     if args.keras_fit and len(args.data_dirs) > 1:
         print("Cannot use keras fit method with multi-genome training.")
         exit()
@@ -198,6 +208,8 @@ def main():
             seqnn_trainer.fit_tape(seqnn_model)
         else:
             seqnn_trainer.fit2(seqnn_model)
+    if args.outlog:
+        sys.stdout = orig_stdout
 
 
 ################################################################################