Enable RNNT ONNX Export #2510

titu1994 · 2021-07-20T02:43:53Z

Changelog

Enable RNNT Export to ONNX (No torchscript yet)
Add inference script for onnx exported rnnt models
Update neural types for export
Refactor Exportable to be less monolithic
Add ExportableEncDecJointModel that is subclass of Exportable

Signed-off-by: smajumdar <titu1994@gmail.com>

* Add parallel block that supports Jasper-like blocks. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Add Encoder with parallel blocks. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Mark test as unit test. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Support pointwise residual connections to connect layers with different fitlers. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Format source code Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Switch to 1D tensor for dropout weights. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Format source code. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com>

Signed-off-by: smajumdar <titu1994@gmail.com>

lgtm-com · 2021-07-20T03:00:53Z

This pull request introduces 2 alerts and fixes 1 when merging 9d626b0 into d08c7a7 - view on LGTM.com

new alerts:

2 for Modification of dictionary returned by locals()

fixed alerts:

1 for Unused import

Signed-off-by: smajumdar <titu1994@gmail.com>

lgtm-com · 2021-07-20T03:19:07Z

This pull request introduces 2 alerts and fixes 1 when merging 2ebd5dd into d08c7a7 - view on LGTM.com

new alerts:

2 for Modification of dictionary returned by locals()

fixed alerts:

1 for Unused import

Signed-off-by: smajumdar <titu1994@gmail.com>

lgtm-com · 2021-07-20T18:08:27Z

This pull request introduces 2 alerts and fixes 1 when merging 1abfbea into d08c7a7 - view on LGTM.com

new alerts:

2 for Modification of dictionary returned by locals()

fixed alerts:

1 for Unused import

borisfom · 2021-07-20T23:23:00Z

nemo/collections/asr/modules/conv_asr.py

@@ -62,8 +63,14 @@ def _prepare_for_export(self, **kwargs):
        m_count = 0
        for m in self.modules():
            if isinstance(m, MaskedConv1d):
-                m.use_mask = False
-                m_count += 1
+                if hasattr(self, '_rnnt_export') and self._rnnt_export:


Why not overriding whole _prepare_for_export for RNNT nets?

That's an option, but that would make framework dependent on what type of encoder was used. This way the Conformer and Citrinet can chose to supply seperate input-output types for CTC and RNNT exports while CTC and RNNT export framework remains unaware of what encoder type it is

Also, I do not see any changes in this PR (nor in main) that would set _rnnt_export attribute - are those part of another PR ?

Part of this PR Inside the asr_model.py file (needs to be expanded)

titu1994 · 2021-07-20T23:39:06Z

nemo/collections/asr/models/asr_model.py

+            format = self.get_format(output)
+
+            # Assign special flag for RNNT export of encoder
+            self.input_module._rnnt_export = True


@borisfom here is where we set it explicitly to the modules. Needs to be before prepare_for_export since samples are already extracted by the time prepare_for_extract is called. (could do with better design but necessary flag to switch from ctc-rnnt onnx export).

Signed-off-by: smajumdar <titu1994@gmail.com>

lgtm-com · 2021-07-21T00:18:44Z

This pull request introduces 2 alerts and fixes 1 when merging 6dff9b0 into c8f9427 - view on LGTM.com

new alerts:

2 for Modification of dictionary returned by locals()

fixed alerts:

1 for Unused import

Signed-off-by: smajumdar <titu1994@gmail.com>

lgtm-com · 2021-07-21T03:20:45Z

This pull request introduces 2 alerts and fixes 1 when merging f6aa380 into c8f9427 - view on LGTM.com

new alerts:

2 for Modification of dictionary returned by locals()

fixed alerts:

1 for Unused import

* Begin export of RNNT models Signed-off-by: smajumdar <titu1994@gmail.com> * Prepare export Signed-off-by: smajumdar <titu1994@gmail.com> * Refactor Exportable Signed-off-by: smajumdar <titu1994@gmail.com> * Enable RNNT export to onnx Signed-off-by: smajumdar <titu1994@gmail.com> * RNNT export Signed-off-by: smajumdar <titu1994@gmail.com> * Attempt to get stateful inference Signed-off-by: smajumdar <titu1994@gmail.com> * Revert changes attempt state fix Signed-off-by: smajumdar <titu1994@gmail.com> * Hack together forward pass <FIX LATER> Signed-off-by: smajumdar <titu1994@gmail.com> * Return length Signed-off-by: smajumdar <titu1994@gmail.com> * Correct signature Signed-off-by: smajumdar <titu1994@gmail.com> * Prepare stateful Signed-off-by: smajumdar <titu1994@gmail.com> * Continie checks for states Signed-off-by: smajumdar <titu1994@gmail.com> * Correct naming Signed-off-by: smajumdar <titu1994@gmail.com> * Correct names Signed-off-by: smajumdar <titu1994@gmail.com> * Pass trace checks Signed-off-by: smajumdar <titu1994@gmail.com> * Update test to be dynamic Signed-off-by: smajumdar <titu1994@gmail.com> * Update naming semantic Signed-off-by: smajumdar <titu1994@gmail.com> * Initial prototype of RNNT greedy decoding Signed-off-by: smajumdar <titu1994@gmail.com> * Fix tracer warnings Signed-off-by: smajumdar <titu1994@gmail.com> * Corrected names of outputs for RNNT encoder Signed-off-by: smajumdar <titu1994@gmail.com> * Fix export rnnt script Signed-off-by: smajumdar <titu1994@gmail.com> * Prototype RNNT forward inference Signed-off-by: smajumdar <titu1994@gmail.com> * Update asr export Signed-off-by: smajumdar <titu1994@gmail.com> * Runnable inference Signed-off-by: smajumdar <titu1994@gmail.com> * Correct log softmax for joint Signed-off-by: smajumdar <titu1994@gmail.com> * Finish RNNT export to onnx Signed-off-by: smajumdar <titu1994@gmail.com> * Fix onnx export Signed-off-by: smajumdar <titu1994@gmail.com> * Create parallel version of ConvASREncoder (NVIDIA#2456) * Add parallel block that supports Jasper-like blocks. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Add Encoder with parallel blocks. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Mark test as unit test. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Support pointwise residual connections to connect layers with different fitlers. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Format source code Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Switch to 1D tensor for dropout weights. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Format source code. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com> * Fix style Signed-off-by: smajumdar <titu1994@gmail.com> * Update comments Signed-off-by: smajumdar <titu1994@gmail.com> * Freeze the nemo model Signed-off-by: smajumdar <titu1994@gmail.com> * Update transcribe script to support RNNT and rnnt model transcriptions Signed-off-by: smajumdar <titu1994@gmail.com> * Add threshold Signed-off-by: smajumdar <titu1994@gmail.com> * Add warning Signed-off-by: smajumdar <titu1994@gmail.com> * Minor fixes Signed-off-by: smajumdar <titu1994@gmail.com> * Style fix Signed-off-by: smajumdar <titu1994@gmail.com> * Correctly replace forward step Signed-off-by: smajumdar <titu1994@gmail.com> * Reset RNNT flag after export is done Signed-off-by: smajumdar <titu1994@gmail.com> * Generalize _rnnt_export flag Signed-off-by: smajumdar <titu1994@gmail.com> * Generalize _rnnt_export flag Signed-off-by: smajumdar <titu1994@gmail.com> Co-authored-by: Aleksei Kalinov <alekseia@nvidia.com>

* Begin export of RNNT models Signed-off-by: smajumdar <titu1994@gmail.com> * Prepare export Signed-off-by: smajumdar <titu1994@gmail.com> * Refactor Exportable Signed-off-by: smajumdar <titu1994@gmail.com> * Enable RNNT export to onnx Signed-off-by: smajumdar <titu1994@gmail.com> * RNNT export Signed-off-by: smajumdar <titu1994@gmail.com> * Attempt to get stateful inference Signed-off-by: smajumdar <titu1994@gmail.com> * Revert changes attempt state fix Signed-off-by: smajumdar <titu1994@gmail.com> * Hack together forward pass <FIX LATER> Signed-off-by: smajumdar <titu1994@gmail.com> * Return length Signed-off-by: smajumdar <titu1994@gmail.com> * Correct signature Signed-off-by: smajumdar <titu1994@gmail.com> * Prepare stateful Signed-off-by: smajumdar <titu1994@gmail.com> * Continie checks for states Signed-off-by: smajumdar <titu1994@gmail.com> * Correct naming Signed-off-by: smajumdar <titu1994@gmail.com> * Correct names Signed-off-by: smajumdar <titu1994@gmail.com> * Pass trace checks Signed-off-by: smajumdar <titu1994@gmail.com> * Update test to be dynamic Signed-off-by: smajumdar <titu1994@gmail.com> * Update naming semantic Signed-off-by: smajumdar <titu1994@gmail.com> * Initial prototype of RNNT greedy decoding Signed-off-by: smajumdar <titu1994@gmail.com> * Fix tracer warnings Signed-off-by: smajumdar <titu1994@gmail.com> * Corrected names of outputs for RNNT encoder Signed-off-by: smajumdar <titu1994@gmail.com> * Fix export rnnt script Signed-off-by: smajumdar <titu1994@gmail.com> * Prototype RNNT forward inference Signed-off-by: smajumdar <titu1994@gmail.com> * Update asr export Signed-off-by: smajumdar <titu1994@gmail.com> * Runnable inference Signed-off-by: smajumdar <titu1994@gmail.com> * Correct log softmax for joint Signed-off-by: smajumdar <titu1994@gmail.com> * Finish RNNT export to onnx Signed-off-by: smajumdar <titu1994@gmail.com> * Fix onnx export Signed-off-by: smajumdar <titu1994@gmail.com> * Create parallel version of ConvASREncoder (NVIDIA#2456) * Add parallel block that supports Jasper-like blocks. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Add Encoder with parallel blocks. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Mark test as unit test. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Support pointwise residual connections to connect layers with different fitlers. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Format source code Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Switch to 1D tensor for dropout weights. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> * Format source code. Signed-off-by: Aleksei Kalinov <alekseia@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com> * Fix style Signed-off-by: smajumdar <titu1994@gmail.com> * Update comments Signed-off-by: smajumdar <titu1994@gmail.com> * Freeze the nemo model Signed-off-by: smajumdar <titu1994@gmail.com> * Update transcribe script to support RNNT and rnnt model transcriptions Signed-off-by: smajumdar <titu1994@gmail.com> * Add threshold Signed-off-by: smajumdar <titu1994@gmail.com> * Add warning Signed-off-by: smajumdar <titu1994@gmail.com> * Minor fixes Signed-off-by: smajumdar <titu1994@gmail.com> * Style fix Signed-off-by: smajumdar <titu1994@gmail.com> * Correctly replace forward step Signed-off-by: smajumdar <titu1994@gmail.com> * Reset RNNT flag after export is done Signed-off-by: smajumdar <titu1994@gmail.com> * Generalize _rnnt_export flag Signed-off-by: smajumdar <titu1994@gmail.com> * Generalize _rnnt_export flag Signed-off-by: smajumdar <titu1994@gmail.com> Co-authored-by: Aleksei Kalinov <alekseia@nvidia.com> Signed-off-by: Paarth Neekhara <paarth.n@gmail.com>

titu1994 and others added 30 commits July 19, 2021 19:42

Begin export of RNNT models

862934a

Signed-off-by: smajumdar <titu1994@gmail.com>

Prepare export

5b6d341

Signed-off-by: smajumdar <titu1994@gmail.com>

Refactor Exportable

1686778

Signed-off-by: smajumdar <titu1994@gmail.com>

Enable RNNT export to onnx

8493e1e

Signed-off-by: smajumdar <titu1994@gmail.com>

RNNT export

f6329bc

Signed-off-by: smajumdar <titu1994@gmail.com>

Attempt to get stateful inference

eb7b20c

Signed-off-by: smajumdar <titu1994@gmail.com>

Revert changes attempt state fix

4a332e6

Signed-off-by: smajumdar <titu1994@gmail.com>

Hack together forward pass <FIX LATER>

508bde9

Signed-off-by: smajumdar <titu1994@gmail.com>

Return length

f601476

Signed-off-by: smajumdar <titu1994@gmail.com>

Correct signature

3c65173

Signed-off-by: smajumdar <titu1994@gmail.com>

Prepare stateful

109916c

Signed-off-by: smajumdar <titu1994@gmail.com>

Continie checks for states

3d1604b

Signed-off-by: smajumdar <titu1994@gmail.com>

Correct naming

b20d8e3

Signed-off-by: smajumdar <titu1994@gmail.com>

Correct names

6de75ec

Signed-off-by: smajumdar <titu1994@gmail.com>

Pass trace checks

eecba8b

Signed-off-by: smajumdar <titu1994@gmail.com>

Update test to be dynamic

0f41dfa

Signed-off-by: smajumdar <titu1994@gmail.com>

Update naming semantic

cefeaa9

Signed-off-by: smajumdar <titu1994@gmail.com>

Initial prototype of RNNT greedy decoding

e98a062

Signed-off-by: smajumdar <titu1994@gmail.com>

Fix tracer warnings

a184817

Signed-off-by: smajumdar <titu1994@gmail.com>

Corrected names of outputs for RNNT encoder

b0cd94e

Signed-off-by: smajumdar <titu1994@gmail.com>

Fix export rnnt script

64c1b06

Signed-off-by: smajumdar <titu1994@gmail.com>

Prototype RNNT forward inference

d065c15

Signed-off-by: smajumdar <titu1994@gmail.com>

Update asr export

3afa684

Signed-off-by: smajumdar <titu1994@gmail.com>

Runnable inference

fec07bb

Signed-off-by: smajumdar <titu1994@gmail.com>

Correct log softmax for joint

b22fbb7

Signed-off-by: smajumdar <titu1994@gmail.com>

Finish RNNT export to onnx

d04497b

Signed-off-by: smajumdar <titu1994@gmail.com>

Fix onnx export

caa08ee

Signed-off-by: smajumdar <titu1994@gmail.com>

Fix style

50d4e8b

Signed-off-by: smajumdar <titu1994@gmail.com>

Update comments

cae154e

Signed-off-by: smajumdar <titu1994@gmail.com>

titu1994 added 4 commits July 19, 2021 19:47

Freeze the nemo model

1bceab9

Signed-off-by: smajumdar <titu1994@gmail.com>

Update transcribe script to support RNNT and rnnt model transcriptions

7812d90

Signed-off-by: smajumdar <titu1994@gmail.com>

Add threshold

571156a

Signed-off-by: smajumdar <titu1994@gmail.com>

Add warning

9d626b0

Signed-off-by: smajumdar <titu1994@gmail.com>

titu1994 force-pushed the rnnt_export branch from a7eeb8f to 9d626b0 Compare July 20, 2021 02:49

Minor fixes

67d9b34

Signed-off-by: smajumdar <titu1994@gmail.com>

titu1994 requested a review from borisfom July 20, 2021 03:04

Style fix

2ebd5dd

Signed-off-by: smajumdar <titu1994@gmail.com>

Correctly replace forward step

1abfbea

Signed-off-by: smajumdar <titu1994@gmail.com>

borisfom reviewed Jul 20, 2021

View reviewed changes

titu1994 commented Jul 20, 2021

View reviewed changes

Reset RNNT flag after export is done

6dff9b0

Signed-off-by: smajumdar <titu1994@gmail.com>

titu1994 added 3 commits July 20, 2021 17:39

Generalize _rnnt_export flag

18bfef0

Signed-off-by: smajumdar <titu1994@gmail.com>

Generalize _rnnt_export flag

4b009d6

Signed-off-by: smajumdar <titu1994@gmail.com>

Merge branch 'main' into rnnt_export

f6aa380

borisfom approved these changes Jul 21, 2021

View reviewed changes

titu1994 merged commit dc69605 into NVIDIA:main Jul 21, 2021

titu1994 deleted the rnnt_export branch July 21, 2021 03:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable RNNT ONNX Export #2510

Enable RNNT ONNX Export #2510

titu1994 commented Jul 20, 2021 •

edited

lgtm-com bot commented Jul 20, 2021

lgtm-com bot commented Jul 20, 2021

lgtm-com bot commented Jul 20, 2021

borisfom Jul 20, 2021 •

edited

titu1994 Jul 20, 2021 •

edited

borisfom Jul 20, 2021

titu1994 Jul 20, 2021

titu1994 Jul 20, 2021 •

edited

lgtm-com bot commented Jul 21, 2021

lgtm-com bot commented Jul 21, 2021

Enable RNNT ONNX Export #2510

Enable RNNT ONNX Export #2510

Conversation

titu1994 commented Jul 20, 2021 • edited

Changelog

lgtm-com bot commented Jul 20, 2021

lgtm-com bot commented Jul 20, 2021

lgtm-com bot commented Jul 20, 2021

borisfom Jul 20, 2021 • edited

Choose a reason for hiding this comment

titu1994 Jul 20, 2021 • edited

Choose a reason for hiding this comment

borisfom Jul 20, 2021

Choose a reason for hiding this comment

titu1994 Jul 20, 2021

Choose a reason for hiding this comment

titu1994 Jul 20, 2021 • edited

Choose a reason for hiding this comment

lgtm-com bot commented Jul 21, 2021

lgtm-com bot commented Jul 21, 2021

titu1994 commented Jul 20, 2021 •

edited

borisfom Jul 20, 2021 •

edited

titu1994 Jul 20, 2021 •

edited

titu1994 Jul 20, 2021 •

edited