Add threads option to prediction (replacing multiprocessing.cpu_count(), set threads=1 for GPU inference).

Set tensorflow version to 2.14.0 for compatibility of using OrderedEnqueuer in training.

Author: XiaotingChen

maxatac/analyses/predict.py

@@ -102,7 +102,7 @@ def run_prediction(args):
                  f"Output filename: {outfile_name_bigwig}"
                  )
 
-    with Pool(int(multiprocessing.cpu_count())) as p:
+    with Pool(args.threads) as p:
         forward_strand_predictions = p.starmap(make_stranded_predictions,
                                                [(regions_pool,
                                                  args.signal,

maxatac/utilities/parser.py

@@ -330,6 +330,12 @@ def get_parser():
                                 help="Skip calling peaks on prediction tracks"
                                 )
 
+    predict_parser.add_argument("--threads",
+                              dest="threads",
+                              type=int,
+                              default=24,
+                              help="Number of processes to run prediction in parallel. Default: 24"
+                              )
     #############################################
     # Train parser
     #############################################

setup.py

@@ -80,7 +80,7 @@ def get_description():
         ]
     ),
     install_requires=[
-        "tensorflow",
+        "tensorflow==2.14.0",
         "tensorboard",
         "biopython",
         "py2bit",