feat(onnx): export ChunkFormer encoder/CTC/RNN-T to ONNX - #42
Merged
Conversation
…n-stream) Add ONNX export + ONNX Runtime inference for ChunkFormer ASR models: - chunkformer/onnx/wrappers.py: ONNX-clean wrappers for full-context and cache-aware streaming encoders, CTC head, RNN-T predictor step and joint. - chunkformer/onnx/runtime.py: OnnxAsrModel host runtime (CTC + RNN-T greedy, streaming chunk loop). - tools/export_onnx.py: per-submodule exporter (CTC/RNN-T, streaming via --streaming, non-stream limited-context via --full-chunk-size/left/right). - tools/verify_onnx_parity.py: PyTorch-vs-ONNX module diff + transcript report. - examples/onnx/README.md, [onnx] extra in pyproject, onnx_out/ gitignored. ONNX-incompatible ops are rewritten only during export (guarded by is_in_onnx_export), leaving eager training/inference unchanged: - attention rel_shift as_strided -> gather; - Tensor.unfold -> gather/unsqueeze (utils.mask.onnx_unfold) in limited-context attention and dynamic conv. Verified parity on chunkformer-ctc-large-vie, chunkformer-rnnt-large-vie and chunkformer-rnnt-small-vie-stream-dct: all module diffs <=2e-5 and transcripts match exactly (full-context, limited-context and streaming).
…s are exact
Add tests/test_onnx_op_rewrites.py covering the two export-only op rewrites:
- onnx_unfold vs native Tensor.unfold: bit-identical across all model call
sites (attn q/kv/mask dims 1/2/-1, dynamic-conv last-dim), 10 shape configs
x {float32, bool, int64}, incl. single-window and trailing-window-drop cases,
plus a check that window count tracks input length.
- rel_shift gather branch vs as_strided branch: exact match over 7 (batch,
head, time1, left, right) configs, plus an explicit index-formula spot check.
39 cases, all passing.
Make an export folder usable for inference without the PyTorch model:
- export_onnx.py now writes vocab.txt and a 'feature' block (sample_rate,
num_mel_bins, frame_length, frame_shift) into the export dir.
- OnnxAsrModel auto-loads vocab.txt (char_dict optional) and adds
extract_features() + transcribe_file() so a new user can run
OnnxAsrModel(dir).transcribe_file('audio.wav') with only the export folder.
Feature deps (torch/torchaudio/pydub) are lazy-imported; transcribe()/encode_*
still need only numpy+onnxruntime. CMVN stays baked in the encoder graph.
- README: simplified self-contained usage + requirements.
Add a true real-time streaming API on top of the exported streaming graphs: - OnnxAsrModel.stream() -> StreamingSession with push_waveform/push_features and finalize(); each call returns the delta text decoded so far. - Manages all streaming caches: encoder attention/conv caches + warm-up offset (CTC and RNN-T), and the RNN-T predictor LSTM state + last token across chunks. - Online fbank extraction with a sample buffer (kaldi snip-edges framing); push_features is a torch-free path (numpy+onnxruntime only). - tools/verify_streaming_session.py: checks online vs offline fbank and that the session transcript matches the offline streaming path (RNN-T + CTC-only). Verified on chunkformer-rnnt-small-vie-stream-dct over audio_1/2/3: online fbank max|diff| 1.9e-6 with identical frame counts, and streaming transcripts match the offline path exactly for both RNN-T and the CTC cache-only path.
- runtime.py: assert optional ORT sessions before .run (union-attr) and annotate fbank/encoder returns (no-any-return). - export_onnx.py: type dynamic nn.Module helper params as Any to silence torch-stub Tensor|Module false positives. black/isort/flake8/mypy all clean on the onnx files.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add ONNX export + ONNX Runtime inference for ChunkFormer ASR models:
ONNX-incompatible ops are rewritten only during export (guarded by is_in_onnx_export), leaving eager training/inference unchanged:
Verified parity on chunkformer-ctc-large-vie, chunkformer-rnnt-large-vie and chunkformer-rnnt-small-vie-stream-dct: all module diffs <=2e-5 and transcripts match exactly (full-context, limited-context and streaming).