Decouple ext scorer init & inference & decoding for the convenience o…

deploy/demo_server.py

@@ -160,22 +160,30 @@ def start_server():
 
     vocab_list = [chars.encode("utf-8") for chars in data_generator.vocab_list]
 
+    if args.decoding_method == "ctc_beam_search":
+        ds2_model.init_ext_scorer(args.alpha, args.beta, args.lang_model_path,
+                                  vocab_list)
     # prepare ASR inference handler
     def file_to_transcript(filename):
         feature = data_generator.process_utterance(filename, "")
-
-        result_transcript = ds2_model.infer_batch(
+        probs_split = ds2_model.infer_batch_probs(
             infer_data=[feature],
-            decoding_method=args.decoding_method,
-            beam_alpha=args.alpha,
-            beam_beta=args.beta,
-            beam_size=args.beam_size,
-            cutoff_prob=args.cutoff_prob,
-            cutoff_top_n=args.cutoff_top_n,
-            vocab_list=vocab_list,
-            language_model_path=args.lang_model_path,
-            num_processes=1,
             feeding_dict=data_generator.feeding)
+
+        if args.decoding_method == "ctc_greedy":
+            result_transcript = ds2_model.decode_batch_greedy(
+                probs_split=probs_split,
+                vocab_list=vocab_list)
+        else:
+            result_transcript = ds2_model.decode_batch_beam_search(
+                probs_split=probs_split,
+                beam_alpha=args.alpha,
+                beam_beta=args.beta,
+                beam_size=args.beam_size,
+                cutoff_prob=args.cutoff_prob,
+                cutoff_top_n=args.cutoff_top_n,
+                vocab_list=vocab_list,
+                num_processes=1)
         return result_transcript[0]
 
     # warming up with utterrances sampled from Librispeech

examples/librispeech/run_tune.sh

@@ -7,7 +7,7 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 \
 python -u tools/tune.py \
 --num_batches=-1 \
 --batch_size=128 \
---trainer_count=8 \
+--trainer_count=4 \
 --beam_size=500 \
 --num_proc_bsearch=12 \
 --num_conv_layers=2 \

infer.py

@@ -90,18 +90,28 @@ def infer():
     # decoders only accept string encoded in utf-8
     vocab_list = [chars.encode("utf-8") for chars in data_generator.vocab_list]
 
-    result_transcripts = ds2_model.infer_batch(
-        infer_data=infer_data,
-        decoding_method=args.decoding_method,
-        beam_alpha=args.alpha,
-        beam_beta=args.beta,
-        beam_size=args.beam_size,
-        cutoff_prob=args.cutoff_prob,
-        cutoff_top_n=args.cutoff_top_n,
-        vocab_list=vocab_list,
-        language_model_path=args.lang_model_path,
-        num_processes=args.num_proc_bsearch,
-        feeding_dict=data_generator.feeding)
+    if args.decoding_method == "ctc_greedy":
+        ds2_model.logger.info("start inference ...")
+        probs_split = ds2_model.infer_batch_probs(infer_data=infer_data,
+            feeding_dict=data_generator.feeding)
+        result_transcripts = ds2_model.decode_batch_greedy(
+            probs_split=probs_split,
+            vocab_list=vocab_list)
+    else:
+        ds2_model.init_ext_scorer(args.alpha, args.beta, args.lang_model_path,
+                                  vocab_list)
+        ds2_model.logger.info("start inference ...")
+        probs_split = ds2_model.infer_batch_probs(infer_data=infer_data,
+            feeding_dict=data_generator.feeding)
+        result_transcripts = ds2_model.decode_batch_beam_search(
+            probs_split=probs_split,
+            beam_alpha=args.alpha,
+            beam_beta=args.beta,
+            beam_size=args.beam_size,
+            cutoff_prob=args.cutoff_prob,
+            cutoff_top_n=args.cutoff_top_n,
+            vocab_list=vocab_list,
+            num_processes=args.num_proc_bsearch)
 
     error_rate_func = cer if args.error_rate_type == 'cer' else wer
     target_transcripts = [data[1] for data in infer_data]

model_utils/model.py

@@ -173,43 +173,19 @@ def infer_loss_batch(self, infer_data):
         # run inference
         return self._loss_inferer.infer(input=infer_data)
 
-    def infer_batch(self, infer_data, decoding_method, beam_alpha, beam_beta,
-                    beam_size, cutoff_prob, cutoff_top_n, vocab_list,
-                    language_model_path, num_processes, feeding_dict):
-        """Model inference. Infer the transcription for a batch of speech
-        utterances.
+    def infer_batch_probs(self, infer_data, feeding_dict):
+        """Infer the prob matrices for a batch of speech utterances.
 
         :param infer_data: List of utterances to infer, with each utterance
                            consisting of a tuple of audio features and
                            transcription text (empty string).
         :type infer_data: list
-        :param decoding_method: Decoding method name, 'ctc_greedy' or
-                                'ctc_beam_search'.
-        :param decoding_method: string
-        :param beam_alpha: Parameter associated with language model.
-        :type beam_alpha: float
-        :param beam_beta: Parameter associated with word count.
-        :type beam_beta: float
-        :param beam_size: Width for Beam search.
-        :type beam_size: int
-        :param cutoff_prob: Cutoff probability in pruning,
-                            default 1.0, no pruning.
-        :type cutoff_prob: float
-        :param cutoff_top_n: Cutoff number in pruning, only top cutoff_top_n
-                        characters with highest probs in vocabulary will be
-                        used in beam search, default 40.
-        :type cutoff_top_n: int
-        :param vocab_list: List of tokens in the vocabulary, for decoding.
-        :type vocab_list: list
-        :param language_model_path: Filepath for language model.
-        :type language_model_path: basestring|None
-        :param num_processes: Number of processes (CPU) for decoder.
-        :type num_processes: int
         :param feeding_dict: Feeding is a map of field name and tuple index
                              of the data that reader returns.
         :type feeding_dict: dict|list
-        :return: List of transcription texts.
-        :rtype: List of basestring
+        :return: List of 2-D probability matrix, and each consists of prob
+                 vectors for one speech utterancce.
+        :rtype: List of matrix
         """
         # define inferer
         if self._inferer == None:
@@ -227,49 +203,102 @@ def infer_batch(self, infer_data, decoding_method, beam_alpha, beam_beta,
             infer_results[start_pos[i]:start_pos[i + 1]]
             for i in xrange(0, len(adapted_infer_data))
         ]
-        # run decoder
+        return probs_split
+
+    def decode_batch_greedy(self, probs_split, vocab_list):
+        """Decode by best path for a batch of probs matrix input.
+
+        :param probs_split: List of 2-D probability matrix, and each consists
+                            of prob vectors for one speech utterancce.
+        :param probs_split: List of matrix
+        :param vocab_list: List of tokens in the vocabulary, for decoding.
+        :type vocab_list: list
+        :return: List of transcription texts.
+        :rtype: List of basestring
+        """
         results = []
-        if decoding_method == "ctc_greedy":
-            # best path decode
-            for i, probs in enumerate(probs_split):
-                output_transcription = ctc_greedy_decoder(
-                    probs_seq=probs, vocabulary=vocab_list)
-                results.append(output_transcription)
-        elif decoding_method == "ctc_beam_search":
-            # initialize external scorer
-            if self._ext_scorer == None:
-                self._loaded_lm_path = language_model_path
-                self.logger.info("begin to initialize the external scorer "
-                                 "for decoding")
-                self._ext_scorer = Scorer(beam_alpha, beam_beta,
-                                          language_model_path, vocab_list)
-
-                lm_char_based = self._ext_scorer.is_character_based()
-                lm_max_order = self._ext_scorer.get_max_order()
-                lm_dict_size = self._ext_scorer.get_dict_size()
-                self.logger.info("language model: "
-                                 "is_character_based = %d," % lm_char_based +
-                                 " max_order = %d," % lm_max_order +
-                                 " dict_size = %d" % lm_dict_size)
-                self.logger.info("end initializing scorer. Start decoding ...")
-            else:
-                self._ext_scorer.reset_params(beam_alpha, beam_beta)
-                assert self._loaded_lm_path == language_model_path
-            # beam search decode
-            num_processes = min(num_processes, len(probs_split))
-            beam_search_results = ctc_beam_search_decoder_batch(
-                probs_split=probs_split,
-                vocabulary=vocab_list,
-                beam_size=beam_size,
-                num_processes=num_processes,
-                ext_scoring_func=self._ext_scorer,
-                cutoff_prob=cutoff_prob,
-                cutoff_top_n=cutoff_top_n)
-
-            results = [result[0][1] for result in beam_search_results]
+        for i, probs in enumerate(probs_split):
+            output_transcription = ctc_greedy_decoder(
+                probs_seq=probs, vocabulary=vocab_list)
+            results.append(output_transcription)
+        return results
+
+    def init_ext_scorer(self, beam_alpha, beam_beta, language_model_path,
+                        vocab_list):
+        """Initialize the external scorer.
+
+        :param beam_alpha: Parameter associated with language model.
+        :type beam_alpha: float
+        :param beam_beta: Parameter associated with word count.
+        :type beam_beta: float
+        :param language_model_path: Filepath for language model. If it is
+                                    empty, the external scorer will be set to
+                                    None, and the decoding method will be pure
+                                    beam search without scorer.
+        :type language_model_path: basestring|None
+        :param vocab_list: List of tokens in the vocabulary, for decoding.
+        :type vocab_list: list
+        """
+        if language_model_path != '':
+            self.logger.info("begin to initialize the external scorer "
+                             "for decoding")
+            self._ext_scorer = Scorer(beam_alpha, beam_beta,
+                                      language_model_path, vocab_list)
+            lm_char_based = self._ext_scorer.is_character_based()
+            lm_max_order = self._ext_scorer.get_max_order()
+            lm_dict_size = self._ext_scorer.get_dict_size()
+            self.logger.info("language model: "
+                             "is_character_based = %d," % lm_char_based +
+                             " max_order = %d," % lm_max_order +
+                             " dict_size = %d" % lm_dict_size)
+            self.logger.info("end initializing scorer")
         else:
-            raise ValueError("Decoding method [%s] is not supported." %
-                             decoding_method)
+            self._ext_scorer = None
+            self.logger.info("no language model provided, "
+                             "decoding by pure beam search without scorer.")
+
+    def decode_batch_beam_search(self, probs_split, beam_alpha, beam_beta,
+                                 beam_size, cutoff_prob, cutoff_top_n,
+                                 vocab_list, num_processes):
+        """Decode by beam search for a batch of probs matrix input.
+
+        :param probs_split: List of 2-D probability matrix, and each consists
+                            of prob vectors for one speech utterancce.
+        :param probs_split: List of matrix
+        :param beam_alpha: Parameter associated with language model.
+        :type beam_alpha: float
+        :param beam_beta: Parameter associated with word count.
+        :type beam_beta: float
+        :param beam_size: Width for Beam search.
+        :type beam_size: int
+        :param cutoff_prob: Cutoff probability in pruning,
+                            default 1.0, no pruning.
+        :type cutoff_prob: float
+        :param cutoff_top_n: Cutoff number in pruning, only top cutoff_top_n
+                        characters with highest probs in vocabulary will be
+                        used in beam search, default 40.
+        :type cutoff_top_n: int
+        :param vocab_list: List of tokens in the vocabulary, for decoding.
+        :type vocab_list: list
+        :param num_processes: Number of processes (CPU) for decoder.
+        :type num_processes: int
+        :return: List of transcription texts.
+        :rtype: List of basestring
+        """
+        if self._ext_scorer != None:
+            self._ext_scorer.reset_params(beam_alpha, beam_beta)
+        # beam search decode
+        num_processes = min(num_processes, len(probs_split))
+        beam_search_results = ctc_beam_search_decoder_batch(
+            probs_split=probs_split,
+            vocabulary=vocab_list,
+            beam_size=beam_size,
+            num_processes=num_processes,
+            ext_scoring_func=self._ext_scorer,
+            cutoff_prob=cutoff_prob,
+            cutoff_top_n=cutoff_top_n)
+
+        results = [result[0][1] for result in beam_search_results]
         return results
 
     def _adapt_feeding_dict(self, feeding_dict):

test.py

@@ -90,22 +90,33 @@ def evaluate():
     # decoders only accept string encoded in utf-8
     vocab_list = [chars.encode("utf-8") for chars in data_generator.vocab_list]
 
+    if args.decoding_method == "ctc_beam_search":
+        ds2_model.init_ext_scorer(args.alpha, args.beta, args.lang_model_path,
+                                  vocab_list)
     errors_func = char_errors if args.error_rate_type == 'cer' else word_errors
     errors_sum, len_refs, num_ins = 0.0, 0, 0
+    ds2_model.logger.info("start evaluation ...")
     for infer_data in batch_reader():
-        result_transcripts = ds2_model.infer_batch(
+        probs_split = ds2_model.infer_batch_probs(
             infer_data=infer_data,
-            decoding_method=args.decoding_method,
-            beam_alpha=args.alpha,
-            beam_beta=args.beta,
-            beam_size=args.beam_size,
-            cutoff_prob=args.cutoff_prob,
-            cutoff_top_n=args.cutoff_top_n,
-            vocab_list=vocab_list,
-            language_model_path=args.lang_model_path,
-            num_processes=args.num_proc_bsearch,
             feeding_dict=data_generator.feeding)
+
+        if args.decoding_method == "ctc_greedy":
+            result_transcripts = ds2_model.decode_batch_greedy(
+                probs_split=probs_split,
+                vocab_list=vocab_list)
+        else:
+            result_transcripts = ds2_model.decode_batch_beam_search(
+                probs_split=probs_split,
+                beam_alpha=args.alpha,
+                beam_beta=args.beta,
+                beam_size=args.beam_size,
+                cutoff_prob=args.cutoff_prob,
+                cutoff_top_n=args.cutoff_top_n,
+                vocab_list=vocab_list,
+                num_processes=args.num_proc_bsearch)
         target_transcripts = [data[1] for data in infer_data]
+
         for target, result in zip(target_transcripts, result_transcripts):
             errors, len_ref = errors_func(target, result)
             errors_sum += errors