Add threads option to prediction (replacing multiprocessing.cpu_count()).

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

Author: XiaotingChen

docs/readme/predict.md

@@ -85,3 +85,7 @@ The windows to use for prediction. These windows must be 1,024 bp wide and have
 ### `-skip_call_peaks, --skip_call_peaks`
 
 This will skip calling peaks at the end of predictions. 
+
+### `--threads`
+
+Set number of parallel threads in prediction tasks. If GPUs are used, set this value to be the number of GPUs used for the task. Default: 24.   

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",