修改解码器输出函数名

Author: yeyupiaoling

export_model.py

@@ -8,7 +8,6 @@
 add_arg = functools.partial(add_arguments, argparser=parser)
 add_arg('configs',          str,   'configs/conformer.yml',    '配置文件')
 add_arg("use_gpu",          bool,  True,                       '是否使用GPU评估模型')
-add_arg("save_quant",       bool,  False,                      '是否保存量化模型')
 add_arg('save_model',       str,   'models/',                  '模型保存的路径')
 add_arg('resume_model',     str,   'models/ConformerModel_fbank/best_model/', '准备转换的模型路径')
 args = parser.parse_args()
@@ -20,5 +19,4 @@
 
 # 导出预测模型
 trainer.export(save_model_path=args.save_model,
-               resume_model=args.resume_model,
-               save_quant=args.save_quant)
+               resume_model=args.resume_model)

masr/decoders/attention_rescoring.py

@@ -61,7 +61,7 @@ def attention_rescoring(
 
         # ctc score in ln domain
         # (beam_size, max_hyps_len, vocab_size)
-        decoder_out, r_decoder_out = model.forward_attention_decoder(hyps_pad, hyps_lens, encoder_out, reverse_weight)
+        decoder_out, r_decoder_out = model.get_decoder_out(hyps_pad, hyps_lens, encoder_out, reverse_weight)
 
         # Only use decoder score for rescoring
         best_score = -float('inf')

masr/model_utils/conformer/model.py

@@ -219,7 +219,7 @@ def get_encoder_out_chunk(self,
         return ctc_probs, att_cache, cnn_cache
 
     @torch.jit.export
-    def forward_attention_decoder(
+    def get_decoder_out(
             self,
             hyps: torch.Tensor,
             hyps_lens: torch.Tensor,

masr/trainer.py

@@ -591,16 +591,13 @@ def evaluate(self, resume_model=None, display_result=False, max_text_duration=No
 
     def export(self,
                save_model_path='models/',
-               resume_model='models/ConformerModel_fbank/best_model/',
-               save_quant=False):
+               resume_model='models/ConformerModel_fbank/best_model/'):
         """导出预测模型
 
         :param save_model_path: 模型保存的路径
         :type save_model_path: str
         :param resume_model: 准备转换的模型路径
         :type resume_model: str
-        :param save_quant: 是否保存量化模型
-        :type save_quant: bool
         :return:
         """
         # 获取训练数据
@@ -632,13 +629,6 @@ def export(self,
         os.makedirs(save_model_dir, exist_ok=True)
         torch.jit.save(infer_model, infer_model_path)
         logger.info("预测模型已保存：{}".format(infer_model_path))
-        # 保存量化模型
-        if save_quant:
-            quant_model_path = os.path.join(os.path.dirname(infer_model_path), 'inference_quant.pth')
-            quantized_model = torch.quantization.quantize_dynamic(self.model)
-            script_quant_model = torch.jit.script(quantized_model)
-            torch.jit.save(script_quant_model, quant_model_path)
-            logger.info("量化模型已保存：{}".format(quant_model_path))
         # 复制词汇表模型
         shutil.copytree(tokenizer.vocab_model_dir, os.path.join(save_model_dir, 'vocab_model'))
         # 保存配置信息
@@ -648,6 +638,9 @@ def export(self,
                 'model_name': self.configs.model_conf.model,
                 'streaming': self.configs.model_conf.model_args.streaming,
                 'sample_rate': self.configs.dataset_conf.dataset.sample_rate,
-                'preprocess_conf': self.configs.preprocess_conf
+                'preprocess_conf': self.configs.preprocess_conf,
             }
+            if self.configs.model_conf.model != "DeepSpeech2Model":
+                inference_config['symbol'] = {'sos': self.model.sos_symbol(), 'eos': self.model.eos_symbol(),
+                                              'ignore_id': self.model.ignore_symbol()}
             json.dump(inference_config, f, indent=4, ensure_ascii=False)