Tensorflow fixes

src/geneml/__init__.py

@@ -1,7 +1,3 @@
-import os
-
 from ._version import __version__
 
 __all__ = ["__version__"]
-
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'  # Suppress TensorFlow logging (1: errors, 2: warnings, 3: info)

src/geneml/main.py

@@ -10,6 +10,7 @@
 
 from geneml.args import parse_args
 from geneml.logger import setup_logger, write_setup_info
+from geneml.model_loader import get_cached_gene_ml_model
 from geneml.outputs import build_cds_sequences, build_prediction_scores_seg, write_fasta, write_gff_file
 from geneml.parallelism import compute_optimal_num_parallelism
 from geneml.params import Params, Strand, build_params_namedtuple
@@ -21,6 +22,18 @@
     run_model,
 )
 
+args = parse_args()
+params = build_params_namedtuple(args)
+
+if params.cpu_only:
+    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+else:
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "1"
+
+# TF needs to be imported after environment variables are set
+import tensorflow as tf # noqa: E402, I001
+
 logger = logging.getLogger("geneml")
 
 
@@ -88,13 +101,6 @@ def process_contig(contig_id: str, seq: str, params: Params, tensorflow_thread_c
     """
     start_time = time.time()
 
-    import tensorflow as tf
-
-    from geneml.model_loader import get_cached_gene_ml_model
-
-    if params.cpu_only:
-        tf.config.set_visible_devices([], 'GPU')
-
     tf.config.threading.set_inter_op_parallelism_threads(tensorflow_thread_count)
     tf.config.threading.set_intra_op_parallelism_threads(tensorflow_thread_count)
 
@@ -262,13 +268,6 @@ def main() -> None:
     Returns:
         None.
     """
-    args = parse_args()
-    params = build_params_namedtuple(args)
-
-    if params.cpu_only:
-        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
-
     logfile = ''.join([params.basepath, '.log'])
     setup_logger(logfile, debug = params.debug, verbose= params.verbose)
     write_setup_info(params)

src/geneml/model_loader.py

@@ -2,7 +2,6 @@
 from functools import cache
 
 import numpy as np
-from keras.src.saving import load_model
 
 MODEL_INTERGENIC = 0
 MODEL_EXON_START = 1
@@ -21,6 +20,7 @@ def __init__(self, path, context_length) -> None:
         """Load a model and initialize context length metadata."""
         self.specify_model_parameters()
 
+        from keras.src.saving import load_model
         self.model = load_model(path, compile=False)
         self.context_length = context_length