| | |
|---|
| | | import logging |
|---|
| | | import torch |
|---|
| | | |
|---|
| | | from funasr.bin.asr_inference_paraformer import Speech2Text |
|---|
| | | from funasr.export.models import get_model |
|---|
| | | import numpy as np |
|---|
| | | import random |
|---|
| | | |
|---|
| | | # torch_version = float(".".join(torch.__version__.split(".")[:2])) |
|---|
| | | # assert torch_version > 1.9 |
|---|
| | | |
|---|
| | | class ASRModelExportParaformer: |
|---|
| | | def __init__(self, cache_dir: Union[Path, str] = None, onnx: bool = True): |
|---|
| | |
|---|
| | | feats_dim=560, |
|---|
| | | onnx=False, |
|---|
| | | ) |
|---|
| | | logging.info("output dir: {}".format(self.cache_dir)) |
|---|
| | | print("output dir: {}".format(self.cache_dir)) |
|---|
| | | self.onnx = onnx |
|---|
| | | |
|---|
| | | |
|---|
| | | def _export( |
|---|
| | | self, |
|---|
| | | model: Speech2Text, |
|---|
| | | model, |
|---|
| | | tag_name: str = None, |
|---|
| | | verbose: bool = False, |
|---|
| | | ): |
|---|
| | |
|---|
| | | model, |
|---|
| | | self.export_config, |
|---|
| | | ) |
|---|
| | | model.eval() |
|---|
| | | # self._export_onnx(model, verbose, export_dir) |
|---|
| | | if self.onnx: |
|---|
| | | self._export_onnx(model, verbose, export_dir) |
|---|
| | | else: |
|---|
| | | self._export_torchscripts(model, verbose, export_dir) |
|---|
| | | |
|---|
| | | logging.info("output dir: {}".format(export_dir)) |
|---|
| | | print("output dir: {}".format(export_dir)) |
|---|
| | | |
|---|
| | | |
|---|
| | | def _export_torchscripts(self, model, verbose, path, enc_size=None): |
|---|
| | | if enc_size: |
|---|
| | | dummy_input = model.get_dummy_inputs(enc_size) |
|---|
| | | else: |
|---|
| | | dummy_input = model.get_dummy_inputs_txt() |
|---|
| | | dummy_input = model.get_dummy_inputs() |
|---|
| | | |
|---|
| | | # model_script = torch.jit.script(model) |
|---|
| | | model_script = torch.jit.trace(model, dummy_input) |
|---|
| | |
|---|
| | | with open(json_file, 'r') as f: |
|---|
| | | config_data = json.load(f) |
|---|
| | | mode = config_data['model']['model_config']['mode'] |
|---|
| | | if mode == 'paraformer': |
|---|
| | | if mode.startswith('paraformer'): |
|---|
| | | from funasr.tasks.asr import ASRTaskParaformer as ASRTask |
|---|
| | | elif mode == 'uniasr': |
|---|
| | | elif mode.startswith('uniasr'): |
|---|
| | | from funasr.tasks.asr import ASRTaskUniASR as ASRTask |
|---|
| | | |
|---|
| | | model, asr_train_args = ASRTask.build_model_from_file( |
|---|
| | |
|---|
| | | dummy_input, |
|---|
| | | os.path.join(path, f'{model.model_name}.onnx'), |
|---|
| | | verbose=verbose, |
|---|
| | | opset_version=12, |
|---|
| | | opset_version=14, |
|---|
| | | input_names=model.get_input_names(), |
|---|
| | | output_names=model.get_output_names(), |
|---|
| | | dynamic_axes=model.get_dynamic_axes() |
|---|
| | | ) |
|---|
| | | |
|---|
| | | |
|---|
| | | if __name__ == '__main__': |
|---|
| | | output_dir = "../export" |
|---|
| | | export_model = ASRModelExportParaformer(cache_dir=output_dir, onnx=True) |
|---|
| | | export_model.export('damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch') |
|---|
| | | import sys |
|---|
| | | |
|---|
| | | model_path = sys.argv[1] |
|---|
| | | output_dir = sys.argv[2] |
|---|
| | | onnx = sys.argv[3] |
|---|
| | | onnx = onnx.lower() |
|---|
| | | onnx = onnx == 'true' |
|---|
| | | # model_path = 'damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch' |
|---|
| | | # output_dir = "../export" |
|---|
| | | export_model = ASRModelExportParaformer(cache_dir=output_dir, onnx=onnx) |
|---|
| | | export_model.export(model_path) |
|---|
| | | # export_model.export('/root/cache/export/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch') |
|---|
