Add option to control number of CPUs in deepcoil standalone. Limit number of cpus during testing

Author: Jan Ludwiczak

bin/deepcoil

@@ -15,11 +15,16 @@ parser.add_argument('-out_path',
                     help='Output directory',
                     default='.',
                     metavar='DIR')
+parser.add_argument('-n_cpu',
+                    help='Number of CPUs to use in the prediction',
+                    default=-1,
+                    type=int,
+                    metavar='NCPU')
 parser.add_argument('--gpu',
-                    help='Use GPU',
+                    help='Use GPU. This option overrides -n_cpu option',
                     action='store_true')
 parser.add_argument('--plot',
-                    help='Plot predictions. Images will be stored in the path defined by -out_path',
+                    help='Plot predictions. Images will be stored in the path defined by the -out_path',
                     action='store_true')
 args = parser.parse_args()
 
@@ -44,7 +49,7 @@ if not len(data) == len(raw_data):
     exit()
 
 print("Loading DeepCoil model...")
-dc = DeepCoil(use_gpu=args.gpu)
+dc = DeepCoil(use_gpu=args.gpu, n_cpu=args.n_cpu)
 
 print('Predicting...')
 preds = dc.predict(data)

tests/test_prediction.py

@@ -12,7 +12,7 @@ def test_seq_regression(self):
         f = h5py.File('tests/data/test_seq_ref.hdf5', 'r')
 
         # Predict with DeepCoil
-        dc = DeepCoil(use_gpu=False)
+        dc = DeepCoil(use_gpu=False, n_cpu=4)
         results = dc.predict(seq_dict)
 
         # Compare predictions with cached predictions