Conformer-Transducer beam search decoding with ONNX models

[OllieBroadhurst]

There is a great ONNXGreedyBatchedRNNTInfer function for performing greedy decoding with .onnx models.

I was wondering how possible it would be to implement something similar for beam search? Specifically for the conformer-transducer which inherently supports beam search.

My initial thought would be to swap out all inference portions of the example script with ONNXRuntime sessions but this lead me down a rabbit hole.

Has this been attempted by anyone yet? I thought I'd check before using up too much of my time taking the wrong approaches.

Thanks!

[titu1994]

Woops, misclicked while typing.

If you look at the exported model inference wrapper - https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L1026-L1196, And compare it with the pytorch greedy batch implementation - https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L479-L814,

You'll see a vast majority of things overlap. With only some numpy or onnx call changes when calling the decoder joint. That's cause the onnx one is a direct copy paste of the pytorch one with just those minor parts modified plus some casts to make everything play nice. Perhaps a few things are hard coded such as the abstract methods of RNNTDecoder since the exported model no longer supports those utility functions

Note that the concrete implementations of actually calling the exported model (say via onnx implementation) are implemented in
https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L1211-L1360 . That should be a generic implementation (mostly) to call the onnx inference runtime.

So you would basically have to follow the same steps for beam search implementations here - https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/parts/submodules/rnnt_beam_decoding.py#L81-L1481 (you probably dont need all 1400 lines if you want to support just one of the beam algorithms)

Do note that this not a trivial effort. Beam search implementation is actually quite a bit more complicated than greedy (plus it doesn't support batching).

[OllieBroadhurst]

Thanks for the pointers, very much appreciated.

[titu1994]

Do note, that the script you mentioned on your original post - that's just a wrapper over these lower level functions. The script will call the python decoding layers, which in turn calls the forward (or call) or one of the decoding strategies - greedy or beam (though in that particular script it's restricted or just beam I think).

Once you implemented the lower level operation in onnx, you can use the script to call your lower level API.

One more (important) thing - I dunno if onnx really has any benefit here. In my naive testing, because I could not implement onnx input binding to GPU, it was a LOT slower than when I implemented the torchscript exported API for inference. So it may be better to stick to TS for the time being + it makes the code simpler

[OllieBroadhurst]

Thanks. I think in this case there could be an issue with exporting to TS? I get the same error as this issue. Looking at the thread, it seems it's not fully supported. Just the case with the encoder, the decoder and joint export fine.

[OllieBroadhurst]

Turns out here's the hack: k2-fsa/sherpa#332 (comment)

[titu1994]

Hmm interesting. @borisfom can we have a flag to disable this optimize for inference step ? ^