[egs] Enable WER eval with GPU (asteroid-team#541)

darius522 · Aug 11, 2021 · 64e10e9 · 64e10e9
1 parent 98d1241
commit 64e10e9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/asteroid/metrics.py b/asteroid/metrics.py
@@ -1,4 +1,5 @@
 import json
+import torch
 import warnings
 import traceback
 from typing import List
@@ -218,17 +219,19 @@ class WERTracker:
         model_name (str): Name of the petrained model to use.
         trans_df (dataframe): Containing field `utt_id` and `text`.
             See librimix/ConvTasNet recipe.
+        use_gpu (bool): Whether to use GPU for forward caculation.
     """
 
-    def __init__(self, model_name, trans_df):
+    def __init__(self, model_name, trans_df, use_gpu=True):
 
         from espnet2.bin.asr_inference import Speech2Text
         from espnet_model_zoo.downloader import ModelDownloader
         import jiwer
 
         self.model_name = model_name
+        self.device = "cuda" if use_gpu else "cpu"
         d = ModelDownloader()
-        self.asr_model = Speech2Text(**d.download_and_unpack(model_name))
+        self.asr_model = Speech2Text(**d.download_and_unpack(model_name), device=self.device)
         self.input_txt_list = []
         self.clean_txt_list = []
         self.output_txt_list = []
@@ -342,6 +345,7 @@ def hsdi(truth, hypothesis, transformation):
         return {k: v for k, v in out if k in keep}
 
     def predict_hypothesis(self, wav):
+        wav = torch.from_numpy(wav).to(self.device)
         nbests = self.asr_model(wav)
         text, *_ = nbests[0]
         return text

diff --git a/egs/librimix/ConvTasNet/eval.py b/egs/librimix/ConvTasNet/eval.py
@@ -71,7 +71,9 @@ def main(conf):
     compute_metrics = update_compute_metrics(conf["compute_wer"], COMPUTE_METRICS)
     anno_df = pd.read_csv(Path(conf["test_dir"]).parent.parent.parent / "test_annotations.csv")
     wer_tracker = (
-        MockWERTracker() if not conf["compute_wer"] else WERTracker(ASR_MODEL_PATH, anno_df)
+        MockWERTracker()
+        if not conf["compute_wer"]
+        else WERTracker(ASR_MODEL_PATH, anno_df, use_gpu=conf["use_gpu"])
     )
     model_path = os.path.join(conf["exp_dir"], "best_model.pth")
     model = ConvTasNet.from_pretrained(model_path)