From b8b764df02827732e2eb8733cff1ce4787f7b246 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B8=B8=E9=9B=81?= Date: Thu, 16 Mar 2023 20:40:21 +0800 Subject: [PATCH] quant --- .../runtime/python/libtorch/torch_paraformer/paraformer_bin.py | 3 +++ .../python/onnxruntime/rapid_paraformer/paraformer_onnx.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/funasr/runtime/python/libtorch/torch_paraformer/paraformer_bin.py b/funasr/runtime/python/libtorch/torch_paraformer/paraformer_bin.py index 3545ccf48..ce975f0af 100644 --- a/funasr/runtime/python/libtorch/torch_paraformer/paraformer_bin.py +++ b/funasr/runtime/python/libtorch/torch_paraformer/paraformer_bin.py @@ -24,12 +24,15 @@ class Paraformer(): device_id: Union[str, int] = "-1", plot_timestamp_to: str = "", pred_bias: int = 1, + quantize: bool = False, ): if not Path(model_dir).exists(): raise FileNotFoundError(f'{model_dir} does not exist.') model_file = os.path.join(model_dir, 'model.torchscripts') + if quantize: + model_file = os.path.join(model_dir, 'model_quant.torchscripts') config_file = os.path.join(model_dir, 'config.yaml') cmvn_file = os.path.join(model_dir, 'am.mvn') config = read_yaml(config_file) diff --git a/funasr/runtime/python/onnxruntime/rapid_paraformer/paraformer_onnx.py b/funasr/runtime/python/onnxruntime/rapid_paraformer/paraformer_onnx.py index 850f007c7..422cb6714 100644 --- a/funasr/runtime/python/onnxruntime/rapid_paraformer/paraformer_onnx.py +++ b/funasr/runtime/python/onnxruntime/rapid_paraformer/paraformer_onnx.py @@ -26,12 +26,15 @@ class Paraformer(): device_id: Union[str, int] = "-1", plot_timestamp_to: str = "", pred_bias: int = 1, + quantize: bool = False, ): if not Path(model_dir).exists(): raise FileNotFoundError(f'{model_dir} does not exist.') model_file = os.path.join(model_dir, 'model.onnx') + if quantize: + model_file = os.path.join(model_dir, 'model_quant.onnx') config_file = os.path.join(model_dir, 'config.yaml') cmvn_file = os.path.join(model_dir, 'am.mvn') config = read_yaml(config_file)