tests/collections/asr/test_asr_exportables.py

# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import tempfile

import onnx
import pytest
import torch.cuda
from omegaconf import DictConfig, ListConfig, OmegaConf

from nemo.collections.asr.models import (
    EncDecClassificationModel,
    EncDecCTCModel,
    EncDecRNNTModel,
    EncDecSpeakerLabelModel,
)
from nemo.collections.asr.parts.utils import asr_module_utils
from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig
from nemo.core.utils import numba_utils
from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__

NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)


class TestExportable:
    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecCTCModel_export_to_onnx(self):
        model_config = DictConfig(
            {
                'preprocessor': DictConfig(self.preprocessor),
                'encoder': DictConfig(self.encoder_dict),
                'decoder': DictConfig(self.decoder_dict),
            }
        )
        model = EncDecCTCModel(cfg=model_config).cuda()
        with tempfile.TemporaryDirectory() as tmpdir:
            filename = os.path.join(tmpdir, 'qn.onnx')
            model.export(
                output=filename, check_trace=True,
            )
            onnx_model = onnx.load(filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.output[0].name == 'logprobs'

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecClassificationModel_export_to_onnx(self, speech_classification_model):
        model = speech_classification_model.cuda()
        with tempfile.TemporaryDirectory() as tmpdir:
            filename = os.path.join(tmpdir, 'edc.onnx')
            model.export(
                output=filename, check_trace=True,
            )
            onnx_model = onnx.load(filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.output[0].name == 'logits'

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecSpeakerLabelModel_export_to_onnx(self, speaker_label_model):
        model = speaker_label_model.cuda()
        with tempfile.TemporaryDirectory() as tmpdir:
            filename = os.path.join(tmpdir, 'sl.onnx')
            model.export(output=filename)
            onnx_model = onnx.load(filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.output[0].name == 'logits'

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecCitrinetModel_export_to_onnx(self, citrinet_model):
        model = citrinet_model.cuda()
        with tempfile.TemporaryDirectory() as tmpdir:
            filename = os.path.join(tmpdir, 'citri.onnx')
            model.export(output=filename)
            onnx_model = onnx.load(filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.input[1].name == 'length'
            assert onnx_model.graph.output[0].name == 'logprobs'

    @pytest.mark.pleasefixme
    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_ConformerModel_export_to_onnx(self, conformer_model):
        model = conformer_model.cuda()
        with tempfile.TemporaryDirectory() as tmpdir, torch.cuda.amp.autocast():
            filename = os.path.join(tmpdir, 'conf.onnx')
            device = next(model.parameters()).device
            input_example = torch.randn(4, model.encoder._feat_in, 777, device=device)
            input_example_length = torch.full(size=(input_example.shape[0],), fill_value=777, device=device)
            model.export(
                output=filename, input_example=tuple([input_example, input_example_length]), check_trace=True,
            )

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_SqueezeformerModel_export_to_onnx(self, squeezeformer_model):
        model = squeezeformer_model.cuda()
        with tempfile.TemporaryDirectory() as tmpdir, torch.cuda.amp.autocast():
            filename = os.path.join(tmpdir, 'squeeze.ts')
            device = next(model.parameters()).device
            input_example = torch.randn(4, model.encoder._feat_in, 777, device=device)
            input_example_length = torch.full(size=(input_example.shape[0],), fill_value=777, device=device)
            model.export(output=filename, input_example=tuple([input_example, input_example_length]), check_trace=True)

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecCitrinetModel_limited_SE_export_to_onnx(self, citrinet_model):
        model = citrinet_model.cuda()
        asr_module_utils.change_conv_asr_se_context_window(model, context_window=24, update_config=False)

        with tempfile.TemporaryDirectory() as tmpdir, torch.cuda.amp.autocast():
            filename = os.path.join(tmpdir, 'citri_se.onnx')
            model.export(
                output=filename, check_trace=True,
            )
            onnx_model = onnx.load(filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.input[1].name == 'length'
            assert onnx_model.graph.output[0].name == 'logprobs'

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecRNNTModel_export_to_onnx(self, citrinet_rnnt_model):
        model = citrinet_rnnt_model.cuda()

        with tempfile.TemporaryDirectory() as tmpdir:
            fn = 'citri_rnnt.onnx'
            filename = os.path.join(tmpdir, fn)
            files, descr = model.export(output=filename, verbose=False)

            encoder_filename = os.path.join(tmpdir, 'encoder-' + fn)
            assert files[0] == encoder_filename
            assert os.path.exists(encoder_filename)
            onnx_model = onnx.load(encoder_filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert len(onnx_model.graph.input) == 2
            assert len(onnx_model.graph.output) == 2
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.input[1].name == 'length'
            assert onnx_model.graph.output[0].name == 'outputs'
            assert onnx_model.graph.output[1].name == 'encoded_lengths'

            decoder_joint_filename = os.path.join(tmpdir, 'decoder_joint-' + fn)
            assert files[1] == decoder_joint_filename
            assert os.path.exists(decoder_joint_filename)
            onnx_model = onnx.load(decoder_joint_filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed

            input_examples = model.decoder.input_example()
            assert type(input_examples[-1]) == tuple
            num_states = len(input_examples[-1])
            state_name = list(model.decoder.output_types.keys())[-1]

            # enc_logits + (all decoder inputs - state tuple) + flattened state list
            assert len(onnx_model.graph.input) == (1 + (len(input_examples) - 1) + num_states)
            assert onnx_model.graph.input[0].name == 'encoder_outputs'
            assert onnx_model.graph.input[1].name == 'targets'
            assert onnx_model.graph.input[2].name == 'target_length'

            if num_states > 0:
                for idx, ip in enumerate(onnx_model.graph.input[3:]):
                    assert ip.name == "input_" + state_name + '_' + str(idx + 1)

            assert len(onnx_model.graph.output) == (len(input_examples) - 1) + num_states
            assert onnx_model.graph.output[0].name == 'outputs'
            assert onnx_model.graph.output[1].name == 'prednet_lengths'

            if num_states > 0:
                for idx, op in enumerate(onnx_model.graph.output[2:]):
                    assert op.name == "output_" + state_name + '_' + str(idx + 1)

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecRNNTModel_export_to_ts(self, citrinet_rnnt_model):
        model = citrinet_rnnt_model.cuda()

        with tempfile.TemporaryDirectory() as tmpdir:
            fn = 'citri_rnnt.ts'
            filename = os.path.join(tmpdir, fn)
            # Perform export + test with the input examples of the RNNT model.
            files, descr = model.export(output=filename, verbose=False, check_trace=True)

            encoder_filename = os.path.join(tmpdir, 'encoder-' + fn)
            assert files[0] == encoder_filename
            assert os.path.exists(encoder_filename)

            ts_encoder = torch.jit.load(encoder_filename)
            assert ts_encoder is not None

            arguments = ts_encoder.forward.schema.arguments[1:]  # First value is `self`
            assert arguments[0].name == 'audio_signal'
            assert arguments[1].name == 'length'

            decoder_joint_filename = os.path.join(tmpdir, 'decoder_joint-' + fn)
            assert files[1] == decoder_joint_filename
            assert os.path.exists(decoder_joint_filename)

            ts_decoder_joint = torch.jit.load(decoder_joint_filename)
            assert ts_decoder_joint is not None

            ts_decoder_joint_args = ts_decoder_joint.forward.schema.arguments[1:]  # First value is self

            input_examples = model.decoder.input_example()
            assert type(input_examples[-1]) == tuple
            num_states = len(input_examples[-1])
            state_name = list(model.decoder.output_types.keys())[-1]

            # enc_logits + (all decoder inputs - state tuple) + flattened state list
            assert len(ts_decoder_joint_args) == (1 + (len(input_examples) - 1) + num_states)
            assert ts_decoder_joint_args[0].name == 'encoder_outputs'
            assert ts_decoder_joint_args[1].name == 'targets'
            assert ts_decoder_joint_args[2].name == 'target_length'

            if num_states > 0:
                for idx, ip in enumerate(ts_decoder_joint_args[3:]):
                    assert ip.name == "input_" + state_name + '_' + str(idx + 1)

    @pytest.mark.run_only_on('GPU')
    @pytest.mark.unit
    def test_EncDecCTCModel_adapted_export_to_onnx(self):
        model_config = DictConfig(
            {
                'preprocessor': DictConfig(self.preprocessor),
                'encoder': DictConfig(self.encoder_dict),
                'decoder': DictConfig(self.decoder_dict),
            }
        )

        # support adapter in encoder
        model_config.encoder.cls = model_config.encoder.cls + 'Adapter'  # ConvASREncoderAdapter

        # load model
        model = EncDecCTCModel(cfg=model_config)

        # add adapter
        adapter_cfg = OmegaConf.structured(
            LinearAdapterConfig(in_features=model_config.encoder.params.jasper[0].filters, dim=32)
        )
        model.add_adapter('temp', cfg=adapter_cfg)

        model = model.cuda()

        with tempfile.TemporaryDirectory() as tmpdir:
            filename = os.path.join(tmpdir, 'qn.onnx')
            model.export(
                output=filename, check_trace=True,
            )
            onnx_model = onnx.load(filename)
            onnx.checker.check_model(onnx_model, full_check=True)  # throws when failed
            assert onnx_model.graph.input[0].name == 'audio_signal'
            assert onnx_model.graph.output[0].name == 'logprobs'

    def setup_method(self):
        self.preprocessor = {
            'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
            'params': dict({}),
        }

        self.encoder_dict = {
            'cls': 'nemo.collections.asr.modules.ConvASREncoder',
            'params': {
                'feat_in': 64,
                'activation': 'relu',
                'conv_mask': True,
                'jasper': [
                    {
                        'filters': 1024,
                        'repeat': 1,
                        'kernel': [1],
                        'stride': [1],
                        'dilation': [1],
                        'dropout': 0.0,
                        'residual': False,
                        'separable': True,
                        'se': True,
                        'se_context_size': -1,
                    }
                ],
            },
        }

        self.decoder_dict = {
            'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
            'params': {
                'feat_in': 1024,
                'num_classes': 28,
                'vocabulary': [
                    ' ',
                    'a',
                    'b',
                    'c',
                    'd',
                    'e',
                    'f',
                    'g',
                    'h',
                    'i',
                    'j',
                    'k',
                    'l',
                    'm',
                    'n',
                    'o',
                    'p',
                    'q',
                    'r',
                    's',
                    't',
                    'u',
                    'v',
                    'w',
                    'x',
                    'y',
                    'z',
                    "'",
                ],
            },
        }


@pytest.fixture()
def speech_classification_model():
    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
    encoder = {
        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
        'params': {
            'feat_in': 64,
            'activation': 'relu',
            'conv_mask': True,
            'jasper': [
                {
                    'filters': 32,
                    'repeat': 1,
                    'kernel': [1],
                    'stride': [1],
                    'dilation': [1],
                    'dropout': 0.0,
                    'residual': False,
                    'separable': True,
                    'se': True,
                    'se_context_size': -1,
                }
            ],
        },
    }

    decoder = {
        'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification',
        'params': {'feat_in': 32, 'num_classes': 30,},
    }

    modelConfig = DictConfig(
        {
            'preprocessor': DictConfig(preprocessor),
            'encoder': DictConfig(encoder),
            'decoder': DictConfig(decoder),
            'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]),
        }
    )
    model = EncDecClassificationModel(cfg=modelConfig)
    return model


@pytest.fixture()
def speaker_label_model():
    preprocessor = {
        '_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor',
    }
    encoder = {
        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
        'feat_in': 64,
        'activation': 'relu',
        'conv_mask': True,
        'jasper': [
            {
                'filters': 512,
                'repeat': 1,
                'kernel': [1],
                'stride': [1],
                'dilation': [1],
                'dropout': 0.0,
                'residual': False,
                'separable': False,
            }
        ],
    }

    decoder = {
        '_target_': 'nemo.collections.asr.modules.SpeakerDecoder',
        'feat_in': 512,
        'num_classes': 2,
        'pool_mode': 'attention',
        'emb_sizes': [1024],
    }

    modelConfig = DictConfig(
        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
    )
    speaker_model = EncDecSpeakerLabelModel(cfg=modelConfig)
    return speaker_model


@pytest.fixture()
def citrinet_model():
    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
    encoder = {
        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
        'params': {
            'feat_in': 80,
            'activation': 'relu',
            'conv_mask': True,
            'jasper': [
                {
                    'filters': 512,
                    'repeat': 1,
                    'kernel': [5],
                    'stride': [1],
                    'dilation': [1],
                    'dropout': 0.0,
                    'residual': False,
                    'separable': True,
                    'se': True,
                    'se_context_size': -1,
                },
                {
                    'filters': 512,
                    'repeat': 5,
                    'kernel': [11],
                    'stride': [2],
                    'dilation': [1],
                    'dropout': 0.1,
                    'residual': True,
                    'separable': True,
                    'se': True,
                    'se_context_size': -1,
                    'stride_last': True,
                    'residual_mode': 'stride_add',
                },
                {
                    'filters': 512,
                    'repeat': 5,
                    'kernel': [13],
                    'stride': [1],
                    'dilation': [1],
                    'dropout': 0.1,
                    'residual': True,
                    'separable': True,
                    'se': True,
                    'se_context_size': -1,
                },
                {
                    'filters': 640,
                    'repeat': 1,
                    'kernel': [41],
                    'stride': [1],
                    'dilation': [1],
                    'dropout': 0.0,
                    'residual': True,
                    'separable': True,
                    'se': True,
                    'se_context_size': -1,
                },
            ],
        },
    }

    decoder = {
        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
        'params': {'feat_in': 640, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))},
    }

    modelConfig = DictConfig(
        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
    )
    citri_model = EncDecCTCModel(cfg=modelConfig)
    return citri_model


@pytest.fixture()
def citrinet_rnnt_model():
    labels = list(chr(i % 28) for i in range(0, 1024))
    model_defaults = {'enc_hidden': 640, 'pred_hidden': 256, 'joint_hidden': 320}

    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
    encoder = {
        '_target_': 'nemo.collections.asr.modules.ConvASREncoder',
        'feat_in': 80,
        'activation': 'relu',
        'conv_mask': True,
        'jasper': [
            {
                'filters': 512,
                'repeat': 1,
                'kernel': [5],
                'stride': [1],
                'dilation': [1],
                'dropout': 0.0,
                'residual': False,
                'separable': True,
                'se': True,
                'se_context_size': -1,
            },
            {
                'filters': 512,
                'repeat': 5,
                'kernel': [11],
                'stride': [2],
                'dilation': [1],
                'dropout': 0.1,
                'residual': True,
                'separable': True,
                'se': True,
                'se_context_size': -1,
                'stride_last': True,
                'residual_mode': 'stride_add',
            },
            {
                'filters': 512,
                'repeat': 5,
                'kernel': [13],
                'stride': [1],
                'dilation': [1],
                'dropout': 0.1,
                'residual': True,
                'separable': True,
                'se': True,
                'se_context_size': -1,
            },
            {
                'filters': 640,
                'repeat': 1,
                'kernel': [41],
                'stride': [1],
                'dilation': [1],
                'dropout': 0.0,
                'residual': True,
                'separable': True,
                'se': True,
                'se_context_size': -1,
            },
        ],
    }

    decoder = {
        '_target_': 'nemo.collections.asr.modules.RNNTDecoder',
        'prednet': {'pred_hidden': 256, 'pred_rnn_layers': 1, 'dropout': 0.0},
    }

    joint = {
        '_target_': 'nemo.collections.asr.modules.RNNTJoint',
        'fuse_loss_wer': False,
        'jointnet': {'joint_hidden': 320, 'activation': 'relu', 'dropout': 0.0},
    }

    decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 5}}

    modelConfig = DictConfig(
        {
            'preprocessor': DictConfig(preprocessor),
            'labels': labels,
            'model_defaults': DictConfig(model_defaults),
            'encoder': DictConfig(encoder),
            'decoder': DictConfig(decoder),
            'joint': DictConfig(joint),
            'decoding': DictConfig(decoding),
        }
    )
    citri_model = EncDecRNNTModel(cfg=modelConfig)
    return citri_model


@pytest.fixture()
def conformer_model():
    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
    encoder = {
        'cls': 'nemo.collections.asr.modules.ConformerEncoder',
        'params': {
            'feat_in': 80,
            'feat_out': -1,
            'n_layers': 2,
            'd_model': 256,
            'subsampling': 'striding',
            'subsampling_factor': 4,
            'subsampling_conv_channels': 512,
            'reduction': None,
            'reduction_position': None,
            'reduction_factor': 1,
            'ff_expansion_factor': 4,
            'self_attention_model': 'rel_pos',
            'n_heads': 8,
            'att_context_size': [-1, -1],
            'xscaling': True,
            'untie_biases': True,
            'pos_emb_max_len': 500,
            'conv_kernel_size': 31,
            'dropout': 0.1,
            'dropout_pre_encoder': 0.1,
            'dropout_emb': 0.0,
            'dropout_att': 0.1,
        },
    }

    decoder = {
        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
        'params': {'feat_in': 256, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))},
    }

    modelConfig = DictConfig(
        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
    )
    conformer_model = EncDecCTCModel(cfg=modelConfig)
    return conformer_model


@pytest.fixture()
def squeezeformer_model():
    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
    encoder = {
        'cls': 'nemo.collections.asr.modules.SqueezeformerEncoder',
        'params': {
            'feat_in': 80,
            'feat_out': -1,
            'n_layers': 2,
            'adaptive_scale': True,
            'time_reduce_idx': 1,
            'time_recovery_idx': None,
            'd_model': 256,
            'subsampling': 'dw_striding',
            'subsampling_factor': 4,
            'subsampling_conv_channels': 512,
            'ff_expansion_factor': 4,
            'self_attention_model': 'rel_pos',
            'n_heads': 8,
            'att_context_size': [-1, -1],
            'xscaling': True,
            'untie_biases': True,
            'pos_emb_max_len': 500,
            'conv_kernel_size': 31,
            'dropout': 0.1,
            'dropout_emb': 0.0,
            'dropout_att': 0.1,
        },
    }

    decoder = {
        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
        'params': {'feat_in': 256, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))},
    }

    modelConfig = DictConfig(
        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
    )
    conformer_model = EncDecCTCModel(cfg=modelConfig)
    return conformer_model