Merge branch 'basemod_gpu' into 'master'

GPU mod base inference enabled

See merge request machine-learning/bonito!83

Author: iiSeymour

README.md

@@ -43,6 +43,8 @@ Modified base calling is handled by [Remora](https://github.com/nanoporetech/rem
 $ bonito basecaller dna_r10.4_e8.1_sup@v3.4 /data/reads --modified-bases 5mC --reference ref.mmi > basecalls_with_mods.bam
 ```
 
+To use the GPU-powered modified bases inference the `onnxruntime-gpu` package is required. 
+
 See available modified base models with the ``remora model list_pretrained`` command.
 
 ## Training your own model

bonito/cli/basecaller.py

@@ -73,7 +73,8 @@ def main(args):
     if args.modified_base_model is not None or args.modified_bases is not None:
         sys.stderr.write("> loading modified base model\n")
         mods_model = load_mods_model(
-            args.modified_bases, args.model_directory, args.modified_base_model
+            args.modified_bases, args.model_directory, args.modified_base_model,
+            device=args.modified_device,
         )
         sys.stderr.write(f"> {mods_model[1]['alphabet_str']}\n")
 
@@ -131,9 +132,12 @@ def main(args):
     )
 
     if mods_model is not None:
-        results = process_itemmap(
-            partial(call_mods, mods_model), results, n_proc=args.modified_procs
-        )
+        if args.modified_device:
+            results = ((k, call_mods(mods_model, k, v)) for k, v in results)
+        else:
+            results = process_itemmap(
+                partial(call_mods, mods_model), results, n_proc=args.modified_procs
+            )
     if aligner:
         results = align_map(aligner, results, n_thread=args.alignment_threads)
 
@@ -167,6 +171,7 @@ def argparser():
     parser.add_argument("--modified-bases", nargs="+")
     parser.add_argument("--modified-base-model")
     parser.add_argument("--modified-procs", default=8, type=int)
+    parser.add_argument("--modified-device", default=None)
     parser.add_argument("--read-ids")
     parser.add_argument("--device", default="cuda")
     parser.add_argument("--seed", default=25, type=int)

bonito/mod_util.py

@@ -31,7 +31,7 @@ def format(self, record):
 log.CONSOLE.setFormatter(CustomFormatter())
 
 
-def load_mods_model(mod_bases, bc_model_str, model_path):
+def load_mods_model(mod_bases, bc_model_str, model_path, device=None):
     if mod_bases is not None:
         try:
             bc_model_type, model_version = bc_model_str.split('@')
@@ -50,8 +50,9 @@ def load_mods_model(mod_bases, bc_model_str, model_path):
             basecall_model_version=model_version,
             modified_bases=mod_bases,
             quiet=True,
+            device=device,
         )
-    return load_model(model_path, quiet=True)
+    return load_model(model_path, quiet=True, device=device)
 
 
 def mods_tags_to_str(mods_tags):