Package: espnet-onnx
espnet_onnxESPnet without PyTorch!Utility library to easily export, quantize, and optimize espnet models to onnx format.
There is no need to install PyTorch or ESPnet on your machine if you already have exported files!espnet_onnx demo in ColabNow demonstration notebook is available in google colab!Simple ASR demo:Simple TTS demo:Installespnet_onnxcan be installed with pippipinstallespnet_onnxIf you want to export pretrained model, you need to installtorch>=1.11.0,espnet,espnet_model_zoo,onnxadditionally.onnx==1.12.0might cause some errors. If you got an error while inference or exporting, please consider downgrading the onnx version.Install guide for developersClone this repository.gitclonegit@github.com:espnet/espnet_onnx.gitCreate virtual environment.cdtools
makevenvexportActivate virtual environment and install torch if required..tools/venv/bin/activate# Please reference official installation guide of PyTorch.pipinstalltorchClone the s3prl repository and install with pip.cdtools
gitclonehttps://github.com/s3prl/s3prlcds3prl
pipinstall.Install warp_transducer for developing transducer model.cdtools
gitclone--single-branch--branchespnet_v1.1https://github.com/b-flo/warp-transducer.gitcdwarp-transducer
mkdirbuild# Please set WITH_OMP to ON or OFF.cdbuild&&cmake-DWITH_OMP="ON"..&&makecdpytorch_binding&&python3-mpipinstall-e.If you want to develop optimization, you also need to develop onnxruntime. Please clone the onnxruntime repository.Since espnet==202308(latest on v0.2.0 release) requiresprotobuf<=3.20.1while the latest onnx requiresprotobuf>=3.20.2, you may get error with installation.
In this case, first, install the espnet==202308, update protobuf to 3.20.3, and then install other libraries.UsageExport modelsespnet_onnxcan export pretrained model published onespnet_model_zoo. By default, exported files will be stored in${HOME}/.cache/espnet_onnx/<tag_name>.fromespnet2.bin.asr_inferenceimportSpeech2Textfromespnet_onnx.exportimportASRModelExportm=ASRModelExport()# download with espnet_model_zoo and export from pretrained modelm.export_from_pretrained('<tag name>',quantize=True)# export from trained modelspeech2text=Speech2Text(args)m.export(speech2text,'<tag name>',quantize=True)You can export pretrained model from zipped file. The zipped file should containmeta.yaml.fromespnet_onnx.exportimportASRModelExportm=ASRModelExport()m.export_from_zip('path/to/the/zipfile',tag_name='tag_name_for_zipped_model',quantize=True)You can set some configuration for export. The available configurations are shown in the details for each models.Details for ASR models and configurationsDetails for TTS models and configurationsDetails for optimization configurationsfromespnet_onnx.exportimportASRModelExportm=ASRModelExport()# Set maximum sequence length to 3000m.set_export_config(max_seq_len=3000)m.export_from_zip('path/to/the/zipfile',tag_name='tag_name_for_zipped_model',)You can easily optimize your model by using theoptimizeoption. If you want to fully optimize your model, you need to install the custom version of onnxruntime fromhere. Please readthis documentfor more detail.fromespnet_onnx.exportimportASRModelExportm=ASRModelExport()m.export_from_zip('path/to/the/zipfile',tag_name='tag_name_for_zipped_model',optimize=True,quantize=True)You can export model from the command line.python-mespnet_onnx.export\--model_typeasr\--input${path_to_zip}\--tagtransformer_lm\--apply_optimize\--apply_quantizeInferenceFor inference,tag_nameormodel_diris used to load onnx file.tag_namehas to be defined intag_config.yamlimportlibrosafromespnet_onnximportSpeech2Textspeech2text=Speech2Text(tag_name='<tag name>')# speech2text = Speech2Text(model_dir='path to the onnx directory')y,sr=librosa.load('sample.wav',sr=16000)nbest=speech2text(y)For streaming asr, you can useStreamingSpeech2Textclass. The speech length should be the same asStreamingSpeech2Text.hop_sizefromespnet_onnximportStreamingSpeech2Textstream_asr=StreamingSpeech2Text(tag_name)# start streaming asrstream_asr.start()whilestreaming:wav=<somecodetogetwav>assertlen(wav)==stream_asr.hop_sizestream_text=stream_asr(wav)[0][0]# You can get non-streaming asr result with end functionnbest=stream_asr.end()You can also simulate streaming model with your wav file withsimulatefunction. PassingTrueas the second argument will show the streaming text as the following code.importlibrosafromespnet_onnximportStreamingSpeech2Textstream_asr=StreamingSpeech2Text(tag_name)y,sr=librosa.load('path/to/wav',sr=16000)nbest=stream_asr.simulate(y,True)# Processing audio with 6 processes.# Result at position 0 :# Result at position 1 :# Result at position 2 : this# Result at position 3 : this is# Result at position 4 : this is a# Result at position 5 : this is aprint(nbest[0][0])# 'this is a pen'If you installed the custom version of onnxruntime, you can run optimized model for inference. You don't have to change any code from the above. If the model was optimized, then espnet_onnx would automatically load the optimized version.You can use only hubert model for your frontend.fromespnet_onnx.exportimportASRModelExport# export your modeltag_name='ESPnet pretrained model with hubert'm=ASRModelExport()m.export_from_pretrained(tag_name,optimize=True)# load only the frontend modelfromespnet_onnx.asr.frontendimportFrontendfrontend=Frontend.get_frontend(tag_name)# use the model in your applicationimportlibrosay,sr=librosa.load('wav file')# y: (B, T)# y_len: (B,)feats=frontend(y[None,:],np.array([len(y)]))If you installedtorchin your environment, you can use frontend in your training.fromespnet_onnx.asr.frontendimportTorchFrontendfrontend=TorchFrontend.get_frontend(tag_name)# load pretrained frontend model# use the model while trainingimportlibrosay,sr=librosa.load('wav file')# You need to place your data on GPU,# and specify the output shape in tupley=torch.Tensor(y).unsqueeze(0).to('cuda')# (1, wav_length)output_shape=(batch_size,feat_length,feats_dims)feats=frontend(y,y.size(1),output_shape)Text2Speech inferenceYou can export TTS models as ASR models.fromespnet2.bin.tts_inferenceimportText2Speechfromespnet_onnx.exportimportTTSModelExportm=TTSModelExport()# download with espnet_model_zoo and export from pretrained modelm.export_from_pretrained('<tag name>',quantize=True)# export from trained modeltext2speech=Text2Speech(args)m.export(text2speech,'<tag name>',quantize=True)You can generate wav files with just simply using the Text2Speech class.fromespnet_onnximportText2Speechtag_name='kan-bayashi/ljspeech_vits'text2speech=Text2Speech(tag_name,use_quantized=True)text='Hello world!'output_dict=text2speech(text)# inference with onnx model.wav=output_dict['wav']How to use GPU on espnet_onnxInstall dependency.First, we needonnxruntime-gpulibrary, instead ofonnxruntime. Please followthis articleto select and install the correct version ofonnxruntime-gpu, depending on your CUDA version.Inference on GPUNow you can speedup the inference speed with GPU. All you need is to select the correct providers, and give it to theSpeech2TextorStreamingSpeech2Textinstance. Seethis articlefor more information about providers.importlibrosafromespnet_onnximportSpeech2TextPROVIDERS=['CUDAExecutionProvider']tag_name='some_tag_name'speech2text=Speech2Text(tag_name,providers=PROVIDERS)y,sr=librosa.load('path/to/wav',sr=16000)nbest=speech2text(y)# runs on GPU.Note that some quantized models are not supported for GPU computation. If you got an error with quantized model, please try not-quantized model.Changes from ESPNetTo avoid the cache problem, I modified some scripts from the original espnet implementation.Add<blank>before<sos>Give sometorch.zeros()arrays to the model.Remove the first token in post process. (removeblank)Replacemake_pad_maskinto new implementation, which can be converted into onnx format.Removedextend_pe()from positional encoding module. The length ofpeis 512 by default.Supported ArchsASR:Supported architecture for ASRTTS:Supported architecture for TTSDeveloper's GuideASR:Developer's GuideReferencesESPNet: end-to-end speech processing toolkitESPNet Model ZooCOPYRIGHTCopyright (c) 2022 Maso SomekiReleased underMIT licenceAuthorMasao Somekicontact:masao.someki@gmail.com