How to use titanet_large.onnx for getting the embeddings of audio file?
State: closed | Pull request: no

Hi everyone, I am trying to use the titanet_large.onnx for getting the embeddings of audio file. Firstly I follow the official instruction to export the onnx version of the pretrained titanet_large model. And then I follow the solution in issue #1825 but it seems not working.

Here is my code
```python
import numpy as np
import pyaudio as pa
import os, time
import librosa
import IPython.display as ipd
import matplotlib.pyplot as plt
import onnxruntime
import nemo
import nemo.collections.asr as nemo_asr

from nemo.core.classes import IterableDataset
from nemo.core.neural_types import NeuralType, AudioSignal, LengthsType
import torch
from torch.utils.data import DataLoader

from omegaconf import OmegaConf
import copy


# simple data layer to pass audio signal
class AudioDataLayer(IterableDataset):
    @property
    def output_types(self):
        return {
            'audio_signal': NeuralType(('B', 'T'), AudioSignal(freq=self._sample_rate)),
            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
        }

    def __init__(self, sample_rate):
        super().__init__()
        self._sample_rate = sample_rate
        self.output = True

    def __iter__(self):
        return self

    def __next__(self):
        if not self.output:
            raise StopIteration
        self.output = False
        return torch.as_tensor(self.signal, dtype=torch.float32), \
            torch.as_tensor(self.signal_shape, dtype=torch.int64)

    def set_signal(self, signal):
        self.signal = signal.astype(np.float32) / 32768.
        self.signal_shape = self.signal.size
        self.output = True

    def __len__(self):
        return 1


# inference method for audio signal (single instance)
def infer_signal(model, signal):
    data_layer.set_signal(signal)
    batch = next(iter(data_loader))
    audio_signal, audio_signal_len = batch
    audio_signal, audio_signal_len = audio_signal.to(model.device), audio_signal_len.to(model.device)
    logits = model.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)
    return logits


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


def infer_signal_onnx(model, signal):
    data_layer.set_signal(signal)
    batch = next(iter(data_loader))
    audio_signal, audio_signal_len = batch
    audio_signal, audio_signal_len = audio_signal.to(model.device), audio_signal_len.to(model.device)
    processed_signal, processed_signal_len = model.preprocessor(
        input_signal=audio_signal, length=audio_signal_len,
    )
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(processed_signal), }
    ologits = ort_session.run(None, ort_inputs)
    alogits = np.asarray(ologits)
    logits = torch.from_numpy(alogits[0])
    return logits


# sample rate, Hz
SAMPLE_RATE = 16000

model = nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(model_name='titanet_large')

# Preserve a copy of the full config
cfg = copy.deepcopy(model._cfg)
print(OmegaConf.to_yaml(cfg))

model.preprocessor = model.from_config_dict(cfg.preprocessor)
model.eval()
model = model.to(model.device)

data_layer = AudioDataLayer(sample_rate=cfg.train_ds.sample_rate)
data_loader = DataLoader(data_layer, batch_size=1, collate_fn=data_layer.collate_fn)

model.export('titanet_l.onnx')
ort_session = onnxruntime.InferenceSession('titanet_l.onnx')

wave_file = 'register.wav'

CHANNELS = 1
RATE = 16000
audio, sample_rate = librosa.load(wave_file, sr=RATE)

s1 = time.time()
infer_signal(model, audio)
print(f"Normal took {time.time()-s1} seconds")

s2 = time.time()
infer_signal_onnx(model, audio)
print(f"Normal took {time.time()-s2} seconds")

```
And below is the error log
```bash
Normal took 0.43860721588134766 seconds
Traceback (most recent call last):
  File "/Users/ericli/mbox-audio-base/test/use_onnx.py", line 112, in <module>
    infer_signal_onnx(model, audio)
  File "/Users/ericli/mbox-audio-base/test/use_onnx.py", line 76, in infer_signal_onnx
    ologits = ort_session.run(None, ort_inputs)
  File "/Users/ericli/miniforge3/envs/audio-base/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 213, in run
    self._validate_input(list(input_feed.keys()))
  File "/Users/ericli/miniforge3/envs/audio-base/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 195, in _validate_input
    raise ValueError(
ValueError: Required inputs (['length']) are missing from input feed (['audio_signal']).
```