From 8d3810e289d44b9e1b86e2f48b6f4f3b3185d62a Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Tue, 7 Feb 2023 15:01:59 +0800
Subject: [PATCH] Simplify ONNX export (#881)

* Simplify ONNX export

* Fix ONNX CI tests
---
 ...pruned-transducer-stateless3-2022-05-13.sh |  30 --
 .github/scripts/test-onnx-export.sh           |  56 +-
 ...runed-transducer-stateless3-2022-05-13.yml |   2 +-
 .../pruned_transducer_stateless2/joiner.py    |   2 -
 .../export-onnx.py                            | 497 ++++++++++++++++++
 .../pruned_transducer_stateless3/export.py    | 270 +---------
 .../onnx_check.py                             | 301 +++++------
 .../onnx_pretrained.py                        | 344 ++++++------
 .../export-onnx.py                            |   2 +-
 .../onnx_check.py                             |   6 -
 .../onnx_pretrained.py                        |   1 -
 11 files changed, 874 insertions(+), 637 deletions(-)
 create mode 100755 egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py

diff --git a/.github/scripts/run-librispeech-pruned-transducer-stateless3-2022-05-13.sh b/.github/scripts/run-librispeech-pruned-transducer-stateless3-2022-05-13.sh
index 880767443..ceb77c7c3 100755
--- a/.github/scripts/run-librispeech-pruned-transducer-stateless3-2022-05-13.sh
+++ b/.github/scripts/run-librispeech-pruned-transducer-stateless3-2022-05-13.sh
@@ -27,14 +27,6 @@ ln -s pretrained-iter-1224000-avg-14.pt pretrained.pt
 ln -s pretrained-iter-1224000-avg-14.pt epoch-99.pt
 popd
 
-log "Test exporting to ONNX format"
-
-./pruned_transducer_stateless3/export.py \
-  --exp-dir $repo/exp \
-  --bpe-model $repo/data/lang_bpe_500/bpe.model \
-  --epoch 99 \
-  --avg 1 \
-  --onnx 1
 
 log "Export to torchscript model"
 ./pruned_transducer_stateless3/export.py \
@@ -51,30 +43,8 @@ log "Export to torchscript model"
   --avg 1 \
   --jit-trace 1
 
-ls -lh $repo/exp/*.onnx
 ls -lh $repo/exp/*.pt
 
-log "Decode with ONNX models"
-
-./pruned_transducer_stateless3/onnx_check.py \
-  --jit-filename $repo/exp/cpu_jit.pt \
-  --onnx-encoder-filename $repo/exp/encoder.onnx \
-  --onnx-decoder-filename $repo/exp/decoder.onnx \
-  --onnx-joiner-filename $repo/exp/joiner.onnx \
-  --onnx-joiner-encoder-proj-filename $repo/exp/joiner_encoder_proj.onnx \
-  --onnx-joiner-decoder-proj-filename $repo/exp/joiner_decoder_proj.onnx
-
-./pruned_transducer_stateless3/onnx_pretrained.py \
-  --bpe-model $repo/data/lang_bpe_500/bpe.model \
-  --encoder-model-filename $repo/exp/encoder.onnx \
-  --decoder-model-filename $repo/exp/decoder.onnx \
-  --joiner-model-filename $repo/exp/joiner.onnx \
-  --joiner-encoder-proj-model-filename $repo/exp/joiner_encoder_proj.onnx \
-  --joiner-decoder-proj-model-filename $repo/exp/joiner_decoder_proj.onnx \
-  $repo/test_wavs/1089-134686-0001.wav \
-  $repo/test_wavs/1221-135766-0001.wav \
-  $repo/test_wavs/1221-135766-0002.wav
-
 log "Decode with models exported by torch.jit.trace()"
 
 ./pruned_transducer_stateless3/jit_pretrained.py \
diff --git a/.github/scripts/test-onnx-export.sh b/.github/scripts/test-onnx-export.sh
index 20aa02950..13a5aa765 100755
--- a/.github/scripts/test-onnx-export.sh
+++ b/.github/scripts/test-onnx-export.sh
@@ -10,9 +10,8 @@ log() {
 
 cd egs/librispeech/ASR
 
-repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-streaming-2022-12-29
-
 log "=========================================================================="
+repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-streaming-2022-12-29
 log "Downloading pre-trained model from $repo_url"
 git lfs install
 GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
@@ -68,3 +67,56 @@ log "Run onnx_pretrained.py"
 
 rm -rf $repo
 log "--------------------------------------------------------------------------"
+
+log "=========================================================================="
+repo_url=https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+log "Downloading pre-trained model from $repo_url"
+git lfs install
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+pushd $repo
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/pretrained-iter-1224000-avg-14.pt"
+
+cd exp
+ln -s pretrained-iter-1224000-avg-14.pt epoch-9999.pt
+popd
+
+log "Export via torch.jit.script()"
+
+./pruned_transducer_stateless3/export.py \
+  --bpe-model $repo/data/lang_bpe_500/bpe.model \
+  --epoch 9999 \
+  --avg 1 \
+  --exp-dir $repo/exp/ \
+  --jit 1
+
+log "Test exporting to ONNX format"
+
+./pruned_transducer_stateless3/export-onnx.py \
+  --bpe-model $repo/data/lang_bpe_500/bpe.model \
+  --epoch 9999 \
+  --avg 1 \
+  --exp-dir $repo/exp/
+
+ls -lh $repo/exp
+
+log "Run onnx_check.py"
+
+./pruned_transducer_stateless3/onnx_check.py \
+  --jit-filename $repo/exp/cpu_jit.pt \
+  --onnx-encoder-filename $repo/exp/encoder-epoch-9999-avg-1.onnx \
+  --onnx-decoder-filename $repo/exp/decoder-epoch-9999-avg-1.onnx \
+  --onnx-joiner-filename $repo/exp/joiner-epoch-9999-avg-1.onnx
+
+log "Run onnx_pretrained.py"
+
+./pruned_transducer_stateless3/onnx_pretrained.py \
+  --encoder-model-filename $repo/exp/encoder-epoch-9999-avg-1.onnx \
+  --decoder-model-filename $repo/exp/decoder-epoch-9999-avg-1.onnx \
+  --joiner-model-filename $repo/exp/joiner-epoch-9999-avg-1.onnx \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
diff --git a/.github/workflows/run-librispeech-pruned-transducer-stateless3-2022-05-13.yml b/.github/workflows/run-librispeech-pruned-transducer-stateless3-2022-05-13.yml
index 2c2bcab0c..f67f7599b 100644
--- a/.github/workflows/run-librispeech-pruned-transducer-stateless3-2022-05-13.yml
+++ b/.github/workflows/run-librispeech-pruned-transducer-stateless3-2022-05-13.yml
@@ -39,7 +39,7 @@ concurrency:
 
 jobs:
   run_librispeech_pruned_transducer_stateless3_2022_05_13:
-    if: github.event.label.name == 'onnx' || github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event_name == 'push' || github.event_name == 'schedule'
+    if: github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event_name == 'push' || github.event_name == 'schedule'
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/joiner.py b/egs/librispeech/ASR/pruned_transducer_stateless2/joiner.py
index 1954f4724..9f88bd029 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless2/joiner.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless2/joiner.py
@@ -56,8 +56,6 @@ class Joiner(nn.Module):
         """
         if not is_jit_tracing():
             assert encoder_out.ndim == decoder_out.ndim
-            assert encoder_out.ndim in (2, 4)
-            assert encoder_out.shape == decoder_out.shape
 
         if project_input:
             logit = self.encoder_proj(encoder_out) + self.decoder_proj(decoder_out)
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py
new file mode 100755
index 000000000..1af68be70
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/export-onnx.py
@@ -0,0 +1,497 @@
+#!/usr/bin/env python3
+#
+# Copyright 2023 Xiaomi Corporation (Author: Fangjun Kuang)
+
+"""
+This script exports a transducer model from PyTorch to ONNX.
+
+We use the pre-trained model from
+https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-streaming-2022-12-29
+as an example to show how to use this file.
+
+1. Download the pre-trained model
+
+cd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+pushd $repo
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/pretrained-iter-1224000-avg-14.pt"
+
+cd exp
+ln -s pretrained-iter-1224000-avg-14.pt epoch-9999.pt
+popd
+
+2. Export the model to ONNX
+
+./pruned_transducer_stateless3/export-onnx.py \
+  --bpe-model $repo/data/lang_bpe_500/bpe.model \
+  --epoch 9999 \
+  --avg 1 \
+  --exp-dir $repo/exp/
+
+It will generate the following 3 files inside $repo/exp:
+
+  - encoder-epoch-9999-avg-1.onnx
+  - decoder-epoch-9999-avg-1.onnx
+  - joiner-epoch-9999-avg-1.onnx
+
+See ./onnx_pretrained.py and ./onnx_check.py for how to
+use the exported ONNX models.
+"""
+
+import argparse
+import logging
+from pathlib import Path
+from typing import Dict, Tuple
+
+import onnx
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from conformer import Conformer
+from decoder import Decoder
+from scaling_converter import convert_scaled_to_non_scaled
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall.checkpoint import average_checkpoints, find_checkpoints, load_checkpoint
+from icefall.utils import setup_logger
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--epoch",
+        type=int,
+        default=28,
+        help="""It specifies the checkpoint to use for averaging.
+        Note: Epoch counts from 0.
+        You can specify --avg to use more checkpoints for model averaging.""",
+    )
+
+    parser.add_argument(
+        "--iter",
+        type=int,
+        default=0,
+        help="""If positive, --epoch is ignored and it
+        will use the checkpoint exp_dir/checkpoint-iter.pt.
+        You can specify --avg to use more checkpoints for model averaging.
+        """,
+    )
+
+    parser.add_argument(
+        "--avg",
+        type=int,
+        default=15,
+        help="Number of checkpoints to average. Automatically select "
+        "consecutive checkpoints before the checkpoint specified by "
+        "'--epoch' and '--iter'",
+    )
+
+    parser.add_argument(
+        "--exp-dir",
+        type=str,
+        default="pruned_transducer_stateless3/exp",
+        help="""It specifies the directory where all training related
+        files, e.g., checkpoints, log, etc, are saved
+        """,
+    )
+
+    parser.add_argument(
+        "--bpe-model",
+        type=str,
+        default="data/lang_bpe_500/bpe.model",
+        help="Path to the BPE model",
+    )
+
+    parser.add_argument(
+        "--context-size",
+        type=int,
+        default=2,
+        help="The context size in the decoder. 1 means bigram; 2 means tri-gram",
+    )
+
+    add_model_arguments(parser)
+
+    return parser
+
+
+def add_meta_data(filename: str, meta_data: Dict[str, str]):
+    """Add meta data to an ONNX model. It is changed in-place.
+
+    Args:
+      filename:
+        Filename of the ONNX model to be changed.
+      meta_data:
+        Key-value pairs.
+    """
+    model = onnx.load(filename)
+    for key, value in meta_data.items():
+        meta = model.metadata_props.add()
+        meta.key = key
+        meta.value = value
+
+    onnx.save(model, filename)
+
+
+class OnnxEncoder(nn.Module):
+    """A wrapper for Conformer and the encoder_proj from the joiner"""
+
+    def __init__(self, encoder: Conformer, encoder_proj: nn.Linear):
+        """
+        Args:
+          encoder:
+            A Conformer encoder.
+          encoder_proj:
+            The projection layer for encoder from the joiner.
+        """
+        super().__init__()
+        self.encoder = encoder
+        self.encoder_proj = encoder_proj
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Please see the help information of Conformer.forward
+
+        Args:
+          x:
+            A 3-D tensor of shape (N, T, C)
+          x_lens:
+            A 1-D tensor of shape (N,). Its dtype is torch.int64
+        Returns:
+          Return a tuple containing:
+            - encoder_out, A 3-D tensor of shape (N, T', joiner_dim)
+            - encoder_out_lens, A 1-D tensor of shape (N,)
+        """
+        encoder_out, encoder_out_lens = self.encoder(x, x_lens)
+
+        encoder_out = self.encoder_proj(encoder_out)
+        # Now encoder_out is of shape (N, T, joiner_dim)
+
+        return encoder_out, encoder_out_lens
+
+
+class OnnxDecoder(nn.Module):
+    """A wrapper for Decoder and the decoder_proj from the joiner"""
+
+    def __init__(self, decoder: Decoder, decoder_proj: nn.Linear):
+        super().__init__()
+        self.decoder = decoder
+        self.decoder_proj = decoder_proj
+
+    def forward(self, y: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+          y:
+            A 2-D tensor of shape (N, context_size).
+        Returns
+          Return a 2-D tensor of shape (N, joiner_dim)
+        """
+        need_pad = False
+        decoder_output = self.decoder(y, need_pad=need_pad)
+        decoder_output = decoder_output.squeeze(1)
+        output = self.decoder_proj(decoder_output)
+
+        return output
+
+
+class OnnxJoiner(nn.Module):
+    """A wrapper for the joiner"""
+
+    def __init__(self, output_linear: nn.Linear):
+        super().__init__()
+        self.output_linear = output_linear
+
+    def forward(
+        self,
+        encoder_out: torch.Tensor,
+        decoder_out: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+          encoder_out:
+            A 2-D tensor of shape (N, joiner_dim)
+          decoder_out:
+            A 2-D tensor of shape (N, joiner_dim)
+        Returns:
+          Return a 2-D tensor of shape (N, vocab_size)
+        """
+        logit = encoder_out + decoder_out
+        logit = self.output_linear(torch.tanh(logit))
+        return logit
+
+
+def export_encoder_model_onnx(
+    encoder_model: OnnxEncoder,
+    encoder_filename: str,
+    opset_version: int = 11,
+) -> None:
+    """Export the given encoder model to ONNX format.
+    The exported model has two inputs:
+
+        - x, a tensor of shape (N, T, C); dtype is torch.float32
+        - x_lens, a tensor of shape (N,); dtype is torch.int64
+
+    and it has two outputs:
+
+        - encoder_out, a tensor of shape (N, T', joiner_dim)
+        - encoder_out_lens, a tensor of shape (N,)
+
+    Args:
+      encoder_model:
+        The input encoder model
+      encoder_filename:
+        The filename to save the exported ONNX model.
+      opset_version:
+        The opset version to use.
+    """
+    x = torch.zeros(1, 100, 80, dtype=torch.float32)
+    x_lens = torch.tensor([100], dtype=torch.int64)
+
+    torch.onnx.export(
+        encoder_model,
+        (x, x_lens),
+        encoder_filename,
+        verbose=False,
+        opset_version=opset_version,
+        input_names=["x", "x_lens"],
+        output_names=["encoder_out", "encoder_out_lens"],
+        dynamic_axes={
+            "x": {0: "N", 1: "T"},
+            "x_lens": {0: "N"},
+            "encoder_out": {0: "N", 1: "T"},
+            "encoder_out_lens": {0: "N"},
+        },
+    )
+
+
+def export_decoder_model_onnx(
+    decoder_model: OnnxDecoder,
+    decoder_filename: str,
+    opset_version: int = 11,
+) -> None:
+    """Export the decoder model to ONNX format.
+
+    The exported model has one input:
+
+        - y: a torch.int64 tensor of shape (N, decoder_model.context_size)
+
+    and has one output:
+
+        - decoder_out: a torch.float32 tensor of shape (N, joiner_dim)
+
+    Args:
+      decoder_model:
+        The decoder model to be exported.
+      decoder_filename:
+        Filename to save the exported ONNX model.
+      opset_version:
+        The opset version to use.
+    """
+    context_size = decoder_model.decoder.context_size
+    vocab_size = decoder_model.decoder.vocab_size
+
+    y = torch.zeros(10, context_size, dtype=torch.int64)
+    torch.onnx.export(
+        decoder_model,
+        y,
+        decoder_filename,
+        verbose=False,
+        opset_version=opset_version,
+        input_names=["y"],
+        output_names=["decoder_out"],
+        dynamic_axes={
+            "y": {0: "N"},
+            "decoder_out": {0: "N"},
+        },
+    )
+
+    meta_data = {
+        "context_size": str(context_size),
+        "vocab_size": str(vocab_size),
+    }
+    add_meta_data(filename=decoder_filename, meta_data=meta_data)
+
+
+def export_joiner_model_onnx(
+    joiner_model: nn.Module,
+    joiner_filename: str,
+    opset_version: int = 11,
+) -> None:
+    """Export the joiner model to ONNX format.
+    The exported joiner model has two inputs:
+
+        - encoder_out: a tensor of shape (N, joiner_dim)
+        - decoder_out: a tensor of shape (N, joiner_dim)
+
+    and produces one output:
+
+        - logit: a tensor of shape (N, vocab_size)
+    """
+    joiner_dim = joiner_model.output_linear.weight.shape[1]
+    logging.info(f"joiner dim: {joiner_dim}")
+
+    projected_encoder_out = torch.rand(11, joiner_dim, dtype=torch.float32)
+    projected_decoder_out = torch.rand(11, joiner_dim, dtype=torch.float32)
+
+    torch.onnx.export(
+        joiner_model,
+        (projected_encoder_out, projected_decoder_out),
+        joiner_filename,
+        verbose=False,
+        opset_version=opset_version,
+        input_names=[
+            "encoder_out",
+            "decoder_out",
+        ],
+        output_names=["logit"],
+        dynamic_axes={
+            "encoder_out": {0: "N"},
+            "decoder_out": {0: "N"},
+            "logit": {0: "N"},
+        },
+    )
+    meta_data = {
+        "joiner_dim": str(joiner_dim),
+    }
+    add_meta_data(filename=joiner_filename, meta_data=meta_data)
+
+
+@torch.no_grad()
+def main():
+    args = get_parser().parse_args()
+    args.exp_dir = Path(args.exp_dir)
+
+    params = get_params()
+    params.update(vars(args))
+
+    device = torch.device("cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+
+    setup_logger(f"{params.exp_dir}/log-export/log-export-onnx")
+
+    logging.info(f"device: {device}")
+
+    sp = spm.SentencePieceProcessor()
+    sp.load(params.bpe_model)
+
+    # <blk> is defined in local/train_bpe_model.py
+    params.blank_id = sp.piece_to_id("<blk>")
+    params.vocab_size = sp.get_piece_size()
+
+    logging.info(params)
+
+    logging.info("About to create model")
+    model = get_transducer_model(params, enable_giga=False)
+
+    model.to(device)
+
+    if params.iter > 0:
+        filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+            : params.avg
+        ]
+        if len(filenames) == 0:
+            raise ValueError(
+                f"No checkpoints found for --iter {params.iter}, --avg {params.avg}"
+            )
+        elif len(filenames) < params.avg:
+            raise ValueError(
+                f"Not enough checkpoints ({len(filenames)}) found for"
+                f" --iter {params.iter}, --avg {params.avg}"
+            )
+        logging.info(f"averaging {filenames}")
+        model.to(device)
+        model.load_state_dict(
+            average_checkpoints(filenames, device=device), strict=False
+        )
+    elif params.avg == 1:
+        load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
+    else:
+        start = params.epoch - params.avg + 1
+        filenames = []
+        for i in range(start, params.epoch + 1):
+            if start >= 0:
+                filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+        logging.info(f"averaging {filenames}")
+        model.to(device)
+        model.load_state_dict(
+            average_checkpoints(filenames, device=device), strict=False
+        )
+
+    model.to("cpu")
+    model.eval()
+
+    convert_scaled_to_non_scaled(model, inplace=True)
+
+    encoder = OnnxEncoder(
+        encoder=model.encoder,
+        encoder_proj=model.joiner.encoder_proj,
+    )
+
+    decoder = OnnxDecoder(
+        decoder=model.decoder,
+        decoder_proj=model.joiner.decoder_proj,
+    )
+
+    joiner = OnnxJoiner(output_linear=model.joiner.output_linear)
+
+    encoder_num_param = sum([p.numel() for p in encoder.parameters()])
+    decoder_num_param = sum([p.numel() for p in decoder.parameters()])
+    joiner_num_param = sum([p.numel() for p in joiner.parameters()])
+    total_num_param = encoder_num_param + decoder_num_param + joiner_num_param
+    logging.info(f"encoder parameters: {encoder_num_param}")
+    logging.info(f"decoder parameters: {decoder_num_param}")
+    logging.info(f"joiner parameters: {joiner_num_param}")
+    logging.info(f"total parameters: {total_num_param}")
+
+    if params.iter > 0:
+        suffix = f"iter-{params.iter}"
+    else:
+        suffix = f"epoch-{params.epoch}"
+
+    suffix += f"-avg-{params.avg}"
+
+    opset_version = 13
+
+    logging.info("Exporting encoder")
+    encoder_filename = params.exp_dir / f"encoder-{suffix}.onnx"
+    export_encoder_model_onnx(
+        encoder,
+        encoder_filename,
+        opset_version=opset_version,
+    )
+    logging.info(f"Exported encoder to {encoder_filename}")
+
+    logging.info("Exporting decoder")
+    decoder_filename = params.exp_dir / f"decoder-{suffix}.onnx"
+    export_decoder_model_onnx(
+        decoder,
+        decoder_filename,
+        opset_version=opset_version,
+    )
+    logging.info(f"Exported decoder to {decoder_filename}")
+
+    logging.info("Exporting joiner")
+    joiner_filename = params.exp_dir / f"joiner-{suffix}.onnx"
+    export_joiner_model_onnx(
+        joiner,
+        joiner_filename,
+        opset_version=opset_version,
+    )
+    logging.info(f"Exported joiner to {joiner_filename}")
+
+
+if __name__ == "__main__":
+    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+
+    main()
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/export.py b/egs/librispeech/ASR/pruned_transducer_stateless3/export.py
index 239bdc12f..f30c9df6a 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/export.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/export.py
@@ -52,32 +52,7 @@ It will also generate 3 other files: `encoder_jit_script.pt`,
 It will generates 3 files: `encoder_jit_trace.pt`,
 `decoder_jit_trace.pt`, and `joiner_jit_trace.pt`.
 
-
-(3) Export to ONNX format
-
-./pruned_transducer_stateless3/export.py \
-  --exp-dir ./pruned_transducer_stateless3/exp \
-  --bpe-model data/lang_bpe_500/bpe.model \
-  --epoch 20 \
-  --avg 10 \
-  --onnx 1
-
-It will generate the following files in the given `exp_dir`.
-Check `onnx_check.py` for how to use them.
-
-    - encoder.onnx
-    - decoder.onnx
-    - joiner.onnx
-    - joiner_encoder_proj.onnx
-    - joiner_decoder_proj.onnx
-
-Please see ./onnx_pretrained.py for usage of the generated files
-
-Check
-https://github.com/k2-fsa/sherpa-onnx
-for how to use the exported models outside of icefall.
-
-(4) Export `model.state_dict()`
+(3) Export `model.state_dict()`
 
 ./pruned_transducer_stateless3/export.py \
   --exp-dir ./pruned_transducer_stateless3/exp \
@@ -210,23 +185,6 @@ def get_parser():
         """,
     )
 
-    parser.add_argument(
-        "--onnx",
-        type=str2bool,
-        default=False,
-        help="""If True, --jit is ignored and it exports the model
-        to onnx format. It will generate the following files:
-
-            - encoder.onnx
-            - decoder.onnx
-            - joiner.onnx
-            - joiner_encoder_proj.onnx
-            - joiner_decoder_proj.onnx
-
-        Refer to ./onnx_check.py and ./onnx_pretrained.py for how to use them.
-        """,
-    )
-
     parser.add_argument(
         "--context-size",
         type=int,
@@ -370,206 +328,6 @@ def export_joiner_model_jit_trace(
     logging.info(f"Saved to {joiner_filename}")
 
 
-def export_encoder_model_onnx(
-    encoder_model: nn.Module,
-    encoder_filename: str,
-    opset_version: int = 11,
-) -> None:
-    """Export the given encoder model to ONNX format.
-    The exported model has two inputs:
-
-        - x, a tensor of shape (N, T, C); dtype is torch.float32
-        - x_lens, a tensor of shape (N,); dtype is torch.int64
-
-    and it has two outputs:
-
-        - encoder_out, a tensor of shape (N, T, C)
-        - encoder_out_lens, a tensor of shape (N,)
-
-    Note: The warmup argument is fixed to 1.
-
-    Args:
-      encoder_model:
-        The input encoder model
-      encoder_filename:
-        The filename to save the exported ONNX model.
-      opset_version:
-        The opset version to use.
-    """
-    x = torch.zeros(1, 100, 80, dtype=torch.float32)
-    x_lens = torch.tensor([100], dtype=torch.int64)
-
-    #  encoder_model = torch.jit.script(encoder_model)
-    # It throws the following error for the above statement
-    #
-    # RuntimeError: Exporting the operator __is_ to ONNX opset version
-    # 11 is not supported. Please feel free to request support or
-    # submit a pull request on PyTorch GitHub.
-    #
-    # I cannot find which statement causes the above error.
-    # torch.onnx.export() will use torch.jit.trace() internally, which
-    # works well for the current reworked model
-    warmup = 1.0
-    torch.onnx.export(
-        encoder_model,
-        (x, x_lens, warmup),
-        encoder_filename,
-        verbose=False,
-        opset_version=opset_version,
-        input_names=["x", "x_lens", "warmup"],
-        output_names=["encoder_out", "encoder_out_lens"],
-        dynamic_axes={
-            "x": {0: "N", 1: "T"},
-            "x_lens": {0: "N"},
-            "encoder_out": {0: "N", 1: "T"},
-            "encoder_out_lens": {0: "N"},
-        },
-    )
-    logging.info(f"Saved to {encoder_filename}")
-
-
-def export_decoder_model_onnx(
-    decoder_model: nn.Module,
-    decoder_filename: str,
-    opset_version: int = 11,
-) -> None:
-    """Export the decoder model to ONNX format.
-
-    The exported model has one input:
-
-        - y: a torch.int64 tensor of shape (N, decoder_model.context_size)
-
-    and has one output:
-
-        - decoder_out: a torch.float32 tensor of shape (N, 1, C)
-
-    Note: The argument need_pad is fixed to False.
-
-    Args:
-      decoder_model:
-        The decoder model to be exported.
-      decoder_filename:
-        Filename to save the exported ONNX model.
-      opset_version:
-        The opset version to use.
-    """
-    y = torch.zeros(10, decoder_model.context_size, dtype=torch.int64)
-    need_pad = False  # Always False, so we can use torch.jit.trace() here
-    # Note(fangjun): torch.jit.trace() is more efficient than torch.jit.script()
-    # in this case
-    torch.onnx.export(
-        decoder_model,
-        (y, need_pad),
-        decoder_filename,
-        verbose=False,
-        opset_version=opset_version,
-        input_names=["y", "need_pad"],
-        output_names=["decoder_out"],
-        dynamic_axes={
-            "y": {0: "N"},
-            "decoder_out": {0: "N"},
-        },
-    )
-    logging.info(f"Saved to {decoder_filename}")
-
-
-def export_joiner_model_onnx(
-    joiner_model: nn.Module,
-    joiner_filename: str,
-    opset_version: int = 11,
-) -> None:
-    """Export the joiner model to ONNX format.
-    The exported joiner model has two inputs:
-
-        - projected_encoder_out: a tensor of shape (N, joiner_dim)
-        - projected_decoder_out: a tensor of shape (N, joiner_dim)
-
-    and produces one output:
-
-        - logit: a tensor of shape (N, vocab_size)
-
-    The exported encoder_proj model has one input:
-
-        - encoder_out: a tensor of shape (N, encoder_out_dim)
-
-    and produces one output:
-
-        - projected_encoder_out: a tensor of shape (N, joiner_dim)
-
-    The exported decoder_proj model has one input:
-
-        - decoder_out: a tensor of shape (N, decoder_out_dim)
-
-    and produces one output:
-
-        - projected_decoder_out: a tensor of shape (N, joiner_dim)
-    """
-    encoder_proj_filename = str(joiner_filename).replace(".onnx", "_encoder_proj.onnx")
-
-    decoder_proj_filename = str(joiner_filename).replace(".onnx", "_decoder_proj.onnx")
-
-    encoder_out_dim = joiner_model.encoder_proj.weight.shape[1]
-    decoder_out_dim = joiner_model.decoder_proj.weight.shape[1]
-    joiner_dim = joiner_model.decoder_proj.weight.shape[0]
-
-    projected_encoder_out = torch.rand(1, joiner_dim, dtype=torch.float32)
-    projected_decoder_out = torch.rand(1, joiner_dim, dtype=torch.float32)
-
-    project_input = False
-    # Note: It uses torch.jit.trace() internally
-    torch.onnx.export(
-        joiner_model,
-        (projected_encoder_out, projected_decoder_out, project_input),
-        joiner_filename,
-        verbose=False,
-        opset_version=opset_version,
-        input_names=[
-            "projected_encoder_out",
-            "projected_decoder_out",
-            "project_input",
-        ],
-        output_names=["logit"],
-        dynamic_axes={
-            "projected_encoder_out": {0: "N"},
-            "projected_decoder_out": {0: "N"},
-            "logit": {0: "N"},
-        },
-    )
-    logging.info(f"Saved to {joiner_filename}")
-
-    encoder_out = torch.rand(1, encoder_out_dim, dtype=torch.float32)
-    torch.onnx.export(
-        joiner_model.encoder_proj,
-        encoder_out,
-        encoder_proj_filename,
-        verbose=False,
-        opset_version=opset_version,
-        input_names=["encoder_out"],
-        output_names=["projected_encoder_out"],
-        dynamic_axes={
-            "encoder_out": {0: "N"},
-            "projected_encoder_out": {0: "N"},
-        },
-    )
-    logging.info(f"Saved to {encoder_proj_filename}")
-
-    decoder_out = torch.rand(1, decoder_out_dim, dtype=torch.float32)
-    torch.onnx.export(
-        joiner_model.decoder_proj,
-        decoder_out,
-        decoder_proj_filename,
-        verbose=False,
-        opset_version=opset_version,
-        input_names=["decoder_out"],
-        output_names=["projected_decoder_out"],
-        dynamic_axes={
-            "decoder_out": {0: "N"},
-            "projected_decoder_out": {0: "N"},
-        },
-    )
-    logging.info(f"Saved to {decoder_proj_filename}")
-
-
 @torch.no_grad()
 def main():
     args = get_parser().parse_args()
@@ -636,31 +394,7 @@ def main():
     model.to("cpu")
     model.eval()
 
-    if params.onnx is True:
-        convert_scaled_to_non_scaled(model, inplace=True)
-        opset_version = 11
-        logging.info("Exporting to onnx format")
-        encoder_filename = params.exp_dir / "encoder.onnx"
-        export_encoder_model_onnx(
-            model.encoder,
-            encoder_filename,
-            opset_version=opset_version,
-        )
-
-        decoder_filename = params.exp_dir / "decoder.onnx"
-        export_decoder_model_onnx(
-            model.decoder,
-            decoder_filename,
-            opset_version=opset_version,
-        )
-
-        joiner_filename = params.exp_dir / "joiner.onnx"
-        export_joiner_model_onnx(
-            model.joiner,
-            joiner_filename,
-            opset_version=opset_version,
-        )
-    elif params.jit is True:
+    if params.jit is True:
         convert_scaled_to_non_scaled(model, inplace=True)
         logging.info("Using torch.jit.script()")
         # We won't use the forward() method of the model in C++, so just ignore
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_check.py b/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_check.py
index 163d737e3..6541f0295 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_check.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_check.py
@@ -19,21 +19,70 @@
 """
 This script checks that exported onnx models produce the same output
 with the given torchscript model for the same input.
+
+We use the pre-trained model from
+https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+as an example to show how to use this file.
+
+1. Download the pre-trained model
+
+cd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+pushd $repo
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/pretrained-iter-1224000-avg-14.pt"
+
+cd exp
+ln -s pretrained-iter-1224000-avg-14.pt epoch-9999.pt
+popd
+
+2. Export the model via torchscript (torch.jit.script())
+
+./pruned_transducer_stateless3/export.py \
+  --bpe-model $repo/data/lang_bpe_500/bpe.model \
+  --epoch 9999 \
+  --avg 1 \
+  --exp-dir $repo/exp/ \
+  --jit 1
+
+It will generate the following file in $repo/exp:
+    - cpu_jit.pt
+
+3. Export the model to ONNX
+
+./pruned_transducer_stateless3/export-onnx.py \
+  --bpe-model $repo/data/lang_bpe_500/bpe.model \
+  --epoch 9999 \
+  --avg 1 \
+  --exp-dir $repo/exp/
+
+It will generate the following 3 files inside $repo/exp:
+
+  - encoder-epoch-9999-avg-1.onnx
+  - decoder-epoch-9999-avg-1.onnx
+  - joiner-epoch-9999-avg-1.onnx
+
+4. Run this file
+
+./pruned_transducer_stateless3/onnx_check.py \
+  --jit-filename $repo/exp/cpu_jit.pt \
+  --onnx-encoder-filename $repo/exp/encoder-epoch-9999-avg-1.onnx \
+  --onnx-decoder-filename $repo/exp/decoder-epoch-9999-avg-1.onnx \
+  --onnx-joiner-filename $repo/exp/joiner-epoch-9999-avg-1.onnx
 """
 
 import argparse
 import logging
 
 from icefall import is_module_available
+from onnx_pretrained import OnnxModel
 
-if not is_module_available("onnxruntime"):
-    raise ValueError("Please 'pip install onnxruntime' first.")
-
-import onnxruntime as ort
 import torch
 
-ort.set_default_logger_severity(3)
-
 
 def get_parser():
     parser = argparse.ArgumentParser(
@@ -68,163 +117,81 @@ def get_parser():
         help="Path to the onnx joiner model",
     )
 
-    parser.add_argument(
-        "--onnx-joiner-encoder-proj-filename",
-        required=True,
-        type=str,
-        help="Path to the onnx joiner encoder projection model",
-    )
-
-    parser.add_argument(
-        "--onnx-joiner-decoder-proj-filename",
-        required=True,
-        type=str,
-        help="Path to the onnx joiner decoder projection model",
-    )
-
     return parser
 
 
 def test_encoder(
-    model: torch.jit.ScriptModule,
-    encoder_session: ort.InferenceSession,
+    torch_model: torch.jit.ScriptModule,
+    onnx_model: OnnxModel,
 ):
-    inputs = encoder_session.get_inputs()
-    outputs = encoder_session.get_outputs()
-    input_names = [n.name for n in inputs]
-    output_names = [n.name for n in outputs]
+    C = 80
+    for i in range(10):
+        N = torch.randint(low=1, high=20, size=(1,)).item()
+        T = torch.randint(low=50, high=100, size=(1,)).item()
+        logging.info(f"test_encoder: iter {i}, N={N}, T={T}")
 
-    assert inputs[0].shape == ["N", "T", 80]
-    assert inputs[1].shape == ["N"]
+        x = torch.rand(N, T, C)
+        x_lens = torch.randint(low=10, high=T + 1, size=(N,))
+        x_lens[0] = T
 
-    for N in [1, 5]:
-        for T in [12, 25]:
-            print("N, T", N, T)
-            x = torch.rand(N, T, 80, dtype=torch.float32)
-            x_lens = torch.randint(low=10, high=T + 1, size=(N,))
-            x_lens[0] = T
+        torch_encoder_out, torch_encoder_out_lens = torch_model.encoder(x, x_lens)
+        torch_encoder_out = torch_model.joiner.encoder_proj(torch_encoder_out)
 
-            encoder_inputs = {
-                input_names[0]: x.numpy(),
-                input_names[1]: x_lens.numpy(),
-            }
-            encoder_out, encoder_out_lens = encoder_session.run(
-                output_names,
-                encoder_inputs,
-            )
+        onnx_encoder_out, onnx_encoder_out_lens = onnx_model.run_encoder(x, x_lens)
 
-            torch_encoder_out, torch_encoder_out_lens = model.encoder(x, x_lens)
-
-            encoder_out = torch.from_numpy(encoder_out)
-            assert torch.allclose(encoder_out, torch_encoder_out, atol=1e-05), (
-                (encoder_out - torch_encoder_out).abs().max(),
-                encoder_out.shape,
-                torch_encoder_out.shape,
-            )
+        assert torch.allclose(torch_encoder_out, onnx_encoder_out, atol=1e-05), (
+            (torch_encoder_out - onnx_encoder_out).abs().max()
+        )
 
 
 def test_decoder(
-    model: torch.jit.ScriptModule,
-    decoder_session: ort.InferenceSession,
+    torch_model: torch.jit.ScriptModule,
+    onnx_model: OnnxModel,
 ):
-    inputs = decoder_session.get_inputs()
-    outputs = decoder_session.get_outputs()
-    input_names = [n.name for n in inputs]
-    output_names = [n.name for n in outputs]
+    context_size = onnx_model.context_size
+    vocab_size = onnx_model.vocab_size
+    for i in range(10):
+        N = torch.randint(1, 100, size=(1,)).item()
+        logging.info(f"test_decoder: iter {i}, N={N}")
+        x = torch.randint(
+            low=1,
+            high=vocab_size,
+            size=(N, context_size),
+            dtype=torch.int64,
+        )
+        torch_decoder_out = torch_model.decoder(x, need_pad=torch.tensor([False]))
+        torch_decoder_out = torch_model.joiner.decoder_proj(torch_decoder_out)
+        torch_decoder_out = torch_decoder_out.squeeze(1)
 
-    assert inputs[0].shape == ["N", 2]
-    for N in [1, 5, 10]:
-        y = torch.randint(low=1, high=500, size=(10, 2))
-
-        decoder_inputs = {input_names[0]: y.numpy()}
-        decoder_out = decoder_session.run(
-            output_names,
-            decoder_inputs,
-        )[0]
-        decoder_out = torch.from_numpy(decoder_out)
-
-        torch_decoder_out = model.decoder(y, need_pad=False)
-        assert torch.allclose(decoder_out, torch_decoder_out, atol=1e-5), (
-            (decoder_out - torch_decoder_out).abs().max()
+        onnx_decoder_out = onnx_model.run_decoder(x)
+        assert torch.allclose(torch_decoder_out, onnx_decoder_out, atol=1e-4), (
+            (torch_decoder_out - onnx_decoder_out).abs().max()
         )
 
 
 def test_joiner(
-    model: torch.jit.ScriptModule,
-    joiner_session: ort.InferenceSession,
-    joiner_encoder_proj_session: ort.InferenceSession,
-    joiner_decoder_proj_session: ort.InferenceSession,
+    torch_model: torch.jit.ScriptModule,
+    onnx_model: OnnxModel,
 ):
-    joiner_inputs = joiner_session.get_inputs()
-    joiner_outputs = joiner_session.get_outputs()
-    joiner_input_names = [n.name for n in joiner_inputs]
-    joiner_output_names = [n.name for n in joiner_outputs]
+    encoder_dim = torch_model.joiner.encoder_proj.weight.shape[1]
+    decoder_dim = torch_model.joiner.decoder_proj.weight.shape[1]
+    for i in range(10):
+        N = torch.randint(1, 100, size=(1,)).item()
+        logging.info(f"test_joiner: iter {i}, N={N}")
+        encoder_out = torch.rand(N, encoder_dim)
+        decoder_out = torch.rand(N, decoder_dim)
 
-    assert joiner_inputs[0].shape == ["N", 512]
-    assert joiner_inputs[1].shape == ["N", 512]
+        projected_encoder_out = torch_model.joiner.encoder_proj(encoder_out)
+        projected_decoder_out = torch_model.joiner.decoder_proj(decoder_out)
 
-    joiner_encoder_proj_inputs = joiner_encoder_proj_session.get_inputs()
-    encoder_proj_input_name = joiner_encoder_proj_inputs[0].name
-
-    assert joiner_encoder_proj_inputs[0].shape == ["N", 512]
-
-    joiner_encoder_proj_outputs = joiner_encoder_proj_session.get_outputs()
-    encoder_proj_output_name = joiner_encoder_proj_outputs[0].name
-
-    joiner_decoder_proj_inputs = joiner_decoder_proj_session.get_inputs()
-    decoder_proj_input_name = joiner_decoder_proj_inputs[0].name
-
-    assert joiner_decoder_proj_inputs[0].shape == ["N", 512]
-
-    joiner_decoder_proj_outputs = joiner_decoder_proj_session.get_outputs()
-    decoder_proj_output_name = joiner_decoder_proj_outputs[0].name
-
-    for N in [1, 5, 10]:
-        encoder_out = torch.rand(N, 512)
-        decoder_out = torch.rand(N, 512)
-
-        projected_encoder_out = torch.rand(N, 512)
-        projected_decoder_out = torch.rand(N, 512)
-
-        joiner_inputs = {
-            joiner_input_names[0]: projected_encoder_out.numpy(),
-            joiner_input_names[1]: projected_decoder_out.numpy(),
-        }
-        joiner_out = joiner_session.run(joiner_output_names, joiner_inputs)[0]
-        joiner_out = torch.from_numpy(joiner_out)
-
-        torch_joiner_out = model.joiner(
-            projected_encoder_out,
-            projected_decoder_out,
-            project_input=False,
-        )
-        assert torch.allclose(joiner_out, torch_joiner_out, atol=1e-5), (
-            (joiner_out - torch_joiner_out).abs().max()
+        torch_joiner_out = torch_model.joiner(encoder_out, decoder_out)
+        onnx_joiner_out = onnx_model.run_joiner(
+            projected_encoder_out, projected_decoder_out
         )
 
-        # Now test encoder_proj
-        joiner_encoder_proj_inputs = {encoder_proj_input_name: encoder_out.numpy()}
-        joiner_encoder_proj_out = joiner_encoder_proj_session.run(
-            [encoder_proj_output_name], joiner_encoder_proj_inputs
-        )[0]
-        joiner_encoder_proj_out = torch.from_numpy(joiner_encoder_proj_out)
-
-        torch_joiner_encoder_proj_out = model.joiner.encoder_proj(encoder_out)
-        assert torch.allclose(
-            joiner_encoder_proj_out, torch_joiner_encoder_proj_out, atol=1e-5
-        ), ((joiner_encoder_proj_out - torch_joiner_encoder_proj_out).abs().max())
-
-        # Now test decoder_proj
-        joiner_decoder_proj_inputs = {decoder_proj_input_name: decoder_out.numpy()}
-        joiner_decoder_proj_out = joiner_decoder_proj_session.run(
-            [decoder_proj_output_name], joiner_decoder_proj_inputs
-        )[0]
-        joiner_decoder_proj_out = torch.from_numpy(joiner_decoder_proj_out)
-
-        torch_joiner_decoder_proj_out = model.joiner.decoder_proj(decoder_out)
-        assert torch.allclose(
-            joiner_decoder_proj_out, torch_joiner_decoder_proj_out, atol=1e-5
-        ), ((joiner_decoder_proj_out - torch_joiner_decoder_proj_out).abs().max())
+        assert torch.allclose(torch_joiner_out, onnx_joiner_out, atol=1e-4), (
+            (torch_joiner_out - onnx_joiner_out).abs().max()
+        )
 
 
 @torch.no_grad()
@@ -232,48 +199,38 @@ def main():
     args = get_parser().parse_args()
     logging.info(vars(args))
 
-    model = torch.jit.load(args.jit_filename)
+    torch_model = torch.jit.load(args.jit_filename)
 
-    options = ort.SessionOptions()
-    options.inter_op_num_threads = 1
-    options.intra_op_num_threads = 1
+    onnx_model = OnnxModel(
+        encoder_model_filename=args.onnx_encoder_filename,
+        decoder_model_filename=args.onnx_decoder_filename,
+        joiner_model_filename=args.onnx_joiner_filename,
+    )
 
     logging.info("Test encoder")
-    encoder_session = ort.InferenceSession(
-        args.onnx_encoder_filename,
-        sess_options=options,
-    )
-    test_encoder(model, encoder_session)
+    test_encoder(torch_model, onnx_model)
 
     logging.info("Test decoder")
-    decoder_session = ort.InferenceSession(
-        args.onnx_decoder_filename,
-        sess_options=options,
-    )
-    test_decoder(model, decoder_session)
+    test_decoder(torch_model, onnx_model)
 
     logging.info("Test joiner")
-    joiner_session = ort.InferenceSession(
-        args.onnx_joiner_filename,
-        sess_options=options,
-    )
-    joiner_encoder_proj_session = ort.InferenceSession(
-        args.onnx_joiner_encoder_proj_filename,
-        sess_options=options,
-    )
-    joiner_decoder_proj_session = ort.InferenceSession(
-        args.onnx_joiner_decoder_proj_filename,
-        sess_options=options,
-    )
-    test_joiner(
-        model,
-        joiner_session,
-        joiner_encoder_proj_session,
-        joiner_decoder_proj_session,
-    )
+    test_joiner(torch_model, onnx_model)
     logging.info("Finished checking ONNX models")
 
 
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+# See https://github.com/pytorch/pytorch/issues/38342
+# and https://github.com/pytorch/pytorch/issues/33354
+#
+# If we don't do this, the delay increases whenever there is
+# a new request that changes the actual batch size.
+# If you use `py-spy dump --pid <server-pid> --native`, you will
+# see a lot of time is spent in re-compiling the torch script model.
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_set_profiling_mode(False)
+torch._C._set_graph_executor_optimize(False)
 if __name__ == "__main__":
     torch.manual_seed(20220727)
     formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_pretrained.py
index 550cf6aad..5adb6c16a 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_pretrained.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/onnx_pretrained.py
@@ -18,35 +18,61 @@
 This script loads ONNX models and uses them to decode waves.
 You can use the following command to get the exported models:
 
-./pruned_transducer_stateless3/export.py \
-  --exp-dir ./pruned_transducer_stateless3/exp \
-  --bpe-model data/lang_bpe_500/bpe.model \
-  --epoch 20 \
-  --avg 10 \
-  --onnx 1
+We use the pre-trained model from
+https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+as an example to show how to use this file.
 
-Usage of this script:
+1. Download the pre-trained model
+
+cd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+pushd $repo
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "exp/pretrained-iter-1224000-avg-14.pt"
+
+cd exp
+ln -s pretrained-iter-1224000-avg-14.pt epoch-9999.pt
+popd
+
+2. Export the model to ONNX
+
+./pruned_transducer_stateless3/export-onnx.py \
+  --bpe-model $repo/data/lang_bpe_500/bpe.model \
+  --epoch 9999 \
+  --avg 1 \
+  --exp-dir $repo/exp/
+
+It will generate the following 3 files inside $repo/exp:
+
+  - encoder-epoch-9999-avg-1.onnx
+  - decoder-epoch-9999-avg-1.onnx
+  - joiner-epoch-9999-avg-1.onnx
+
+3. Run this file
 
 ./pruned_transducer_stateless3/onnx_pretrained.py \
-  --encoder-model-filename ./pruned_transducer_stateless3/exp/encoder.onnx \
-  --decoder-model-filename ./pruned_transducer_stateless3/exp/decoder.onnx \
-  --joiner-model-filename ./pruned_transducer_stateless3/exp/joiner.onnx \
-  --joiner-encoder-proj-model-filename ./pruned_transducer_stateless3/exp/joiner_encoder_proj.onnx \
-  --joiner-decoder-proj-model-filename ./pruned_transducer_stateless3/exp/joiner_decoder_proj.onnx \
-  --bpe-model ./data/lang_bpe_500/bpe.model \
-  /path/to/foo.wav \
-  /path/to/bar.wav
+  --encoder-model-filename $repo/exp/encoder-epoch-9999-avg-1.onnx \
+  --decoder-model-filename $repo/exp/decoder-epoch-9999-avg-1.onnx \
+  --joiner-model-filename $repo/exp/joiner-epoch-9999-avg-1.onnx \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
 """
 
 import argparse
 import logging
 import math
-from typing import List
+from typing import List, Tuple
 
+import k2
 import kaldifeat
 import numpy as np
 import onnxruntime as ort
-import sentencepiece as spm
 import torch
 import torchaudio
 from torch.nn.utils.rnn import pad_sequence
@@ -79,23 +105,9 @@ def get_parser():
     )
 
     parser.add_argument(
-        "--joiner-encoder-proj-model-filename",
+        "--tokens",
         type=str,
-        required=True,
-        help="Path to the joiner encoder_proj onnx model. ",
-    )
-
-    parser.add_argument(
-        "--joiner-decoder-proj-model-filename",
-        type=str,
-        required=True,
-        help="Path to the joiner decoder_proj onnx model. ",
-    )
-
-    parser.add_argument(
-        "--bpe-model",
-        type=str,
-        help="""Path to bpe.model.""",
+        help="""Path to tokens.txt.""",
     )
 
     parser.add_argument(
@@ -115,16 +127,122 @@ def get_parser():
         help="The sample rate of the input sound file",
     )
 
-    parser.add_argument(
-        "--context-size",
-        type=int,
-        default=2,
-        help="Context size of the decoder model",
-    )
-
     return parser
 
 
+class OnnxModel:
+    def __init__(
+        self,
+        encoder_model_filename: str,
+        decoder_model_filename: str,
+        joiner_model_filename: str,
+    ):
+        session_opts = ort.SessionOptions()
+        session_opts.inter_op_num_threads = 1
+        session_opts.intra_op_num_threads = 1
+
+        self.session_opts = session_opts
+
+        self.init_encoder(encoder_model_filename)
+        self.init_decoder(decoder_model_filename)
+        self.init_joiner(joiner_model_filename)
+
+    def init_encoder(self, encoder_model_filename: str):
+        self.encoder = ort.InferenceSession(
+            encoder_model_filename,
+            sess_options=self.session_opts,
+        )
+
+    def init_decoder(self, decoder_model_filename: str):
+        self.decoder = ort.InferenceSession(
+            decoder_model_filename,
+            sess_options=self.session_opts,
+        )
+
+        decoder_meta = self.decoder.get_modelmeta().custom_metadata_map
+        self.context_size = int(decoder_meta["context_size"])
+        self.vocab_size = int(decoder_meta["vocab_size"])
+
+        logging.info(f"context_size: {self.context_size}")
+        logging.info(f"vocab_size: {self.vocab_size}")
+
+    def init_joiner(self, joiner_model_filename: str):
+        self.joiner = ort.InferenceSession(
+            joiner_model_filename,
+            sess_options=self.session_opts,
+        )
+
+        joiner_meta = self.joiner.get_modelmeta().custom_metadata_map
+        self.joiner_dim = int(joiner_meta["joiner_dim"])
+
+        logging.info(f"joiner_dim: {self.joiner_dim}")
+
+    def run_encoder(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+          x:
+            A 3-D tensor of shape (N, T, C)
+          x_lens:
+            A 2-D tensor of shape (N,). Its dtype is torch.int64
+        Returns:
+          Return a tuple containing:
+            - encoder_out, its shape is (N, T', joiner_dim)
+            - encoder_out_lens, its shape is (N,)
+        """
+        out = self.encoder.run(
+            [
+                self.encoder.get_outputs()[0].name,
+                self.encoder.get_outputs()[1].name,
+            ],
+            {
+                self.encoder.get_inputs()[0].name: x.numpy(),
+                self.encoder.get_inputs()[1].name: x_lens.numpy(),
+            },
+        )
+        return torch.from_numpy(out[0]), torch.from_numpy(out[1])
+
+    def run_decoder(self, decoder_input: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+          decoder_input:
+            A 2-D tensor of shape (N, context_size)
+        Returns:
+          Return a 2-D tensor of shape (N, joiner_dim)
+        """
+        out = self.decoder.run(
+            [self.decoder.get_outputs()[0].name],
+            {self.decoder.get_inputs()[0].name: decoder_input.numpy()},
+        )[0]
+
+        return torch.from_numpy(out)
+
+    def run_joiner(
+        self, encoder_out: torch.Tensor, decoder_out: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+          encoder_out:
+            A 2-D tensor of shape (N, joiner_dim)
+          decoder_out:
+            A 2-D tensor of shape (N, joiner_dim)
+        Returns:
+          Return a 2-D tensor of shape (N, vocab_size)
+        """
+        out = self.joiner.run(
+            [self.joiner.get_outputs()[0].name],
+            {
+                self.joiner.get_inputs()[0].name: encoder_out.numpy(),
+                self.joiner.get_inputs()[1].name: decoder_out.numpy(),
+            },
+        )[0]
+
+        return torch.from_numpy(out)
+
+
 def read_sound_files(
     filenames: List[str], expected_sample_rate: float
 ) -> List[torch.Tensor]:
@@ -149,36 +267,22 @@ def read_sound_files(
 
 
 def greedy_search(
-    decoder: ort.InferenceSession,
-    joiner: ort.InferenceSession,
-    joiner_encoder_proj: ort.InferenceSession,
-    joiner_decoder_proj: ort.InferenceSession,
-    encoder_out: np.ndarray,
-    encoder_out_lens: np.ndarray,
-    context_size: int,
+    model: OnnxModel,
+    encoder_out: torch.Tensor,
+    encoder_out_lens: torch.Tensor,
 ) -> List[List[int]]:
     """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
     Args:
-      decoder:
-        The decoder model.
-      joiner:
-        The joiner model.
-      joiner_encoder_proj:
-        The joiner encoder projection model.
-      joiner_decoder_proj:
-        The joiner decoder projection model.
+      model:
+        The transducer model.
       encoder_out:
-        A 3-D tensor of shape (N, T, C)
+        A 3-D tensor of shape (N, T, joiner_dim)
       encoder_out_lens:
         A 1-D tensor of shape (N,).
-      context_size:
-        The context size of the decoder model.
     Returns:
       Return the decoded results for each utterance.
     """
-    encoder_out = torch.from_numpy(encoder_out)
-    encoder_out_lens = torch.from_numpy(encoder_out_lens)
-    assert encoder_out.ndim == 3
+    assert encoder_out.ndim == 3, encoder_out.shape
     assert encoder_out.size(0) >= 1, encoder_out.size(0)
 
     packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
@@ -188,11 +292,6 @@ def greedy_search(
         enforce_sorted=False,
     )
 
-    projected_encoder_out = joiner_encoder_proj.run(
-        [joiner_encoder_proj.get_outputs()[0].name],
-        {joiner_encoder_proj.get_inputs()[0].name: packed_encoder_out.data.numpy()},
-    )[0]
-
     blank_id = 0  # hard-code to 0
 
     batch_size_list = packed_encoder_out.batch_sizes.tolist()
@@ -201,50 +300,27 @@ def greedy_search(
     assert torch.all(encoder_out_lens > 0), encoder_out_lens
     assert N == batch_size_list[0], (N, batch_size_list)
 
+    context_size = model.context_size
     hyps = [[blank_id] * context_size for _ in range(N)]
 
-    decoder_input_nodes = decoder.get_inputs()
-    decoder_output_nodes = decoder.get_outputs()
-
-    joiner_input_nodes = joiner.get_inputs()
-    joiner_output_nodes = joiner.get_outputs()
-
     decoder_input = torch.tensor(
         hyps,
         dtype=torch.int64,
     )  # (N, context_size)
 
-    decoder_out = decoder.run(
-        [decoder_output_nodes[0].name],
-        {
-            decoder_input_nodes[0].name: decoder_input.numpy(),
-        },
-    )[0].squeeze(1)
-    projected_decoder_out = joiner_decoder_proj.run(
-        [joiner_decoder_proj.get_outputs()[0].name],
-        {joiner_decoder_proj.get_inputs()[0].name: decoder_out},
-    )[0]
-
-    projected_decoder_out = torch.from_numpy(projected_decoder_out)
+    decoder_out = model.run_decoder(decoder_input)
 
     offset = 0
     for batch_size in batch_size_list:
         start = offset
         end = offset + batch_size
-        current_encoder_out = projected_encoder_out[start:end]
-        # current_encoder_out's shape: (batch_size, encoder_out_dim)
+        current_encoder_out = packed_encoder_out.data[start:end]
+        # current_encoder_out's shape: (batch_size, joiner_dim)
         offset = end
 
-        projected_decoder_out = projected_decoder_out[:batch_size]
+        decoder_out = decoder_out[:batch_size]
+        logits = model.run_joiner(current_encoder_out, decoder_out)
 
-        logits = joiner.run(
-            [joiner_output_nodes[0].name],
-            {
-                joiner_input_nodes[0].name: current_encoder_out,
-                joiner_input_nodes[1].name: projected_decoder_out.numpy(),
-            },
-        )[0]
-        logits = torch.from_numpy(logits).squeeze(1).squeeze(1)
         # logits'shape (batch_size, vocab_size)
 
         assert logits.ndim == 2, logits.shape
@@ -261,17 +337,7 @@ def greedy_search(
                 decoder_input,
                 dtype=torch.int64,
             )
-            decoder_out = decoder.run(
-                [decoder_output_nodes[0].name],
-                {
-                    decoder_input_nodes[0].name: decoder_input.numpy(),
-                },
-            )[0].squeeze(1)
-            projected_decoder_out = joiner_decoder_proj.run(
-                [joiner_decoder_proj.get_outputs()[0].name],
-                {joiner_decoder_proj.get_inputs()[0].name: decoder_out},
-            )[0]
-            projected_decoder_out = torch.from_numpy(projected_decoder_out)
+            decoder_out = model.run_decoder(decoder_input)
 
     sorted_ans = [h[context_size:] for h in hyps]
     ans = []
@@ -287,39 +353,12 @@ def main():
     parser = get_parser()
     args = parser.parse_args()
     logging.info(vars(args))
-
-    session_opts = ort.SessionOptions()
-    session_opts.inter_op_num_threads = 1
-    session_opts.intra_op_num_threads = 1
-
-    encoder = ort.InferenceSession(
-        args.encoder_model_filename,
-        sess_options=session_opts,
+    model = OnnxModel(
+        encoder_model_filename=args.encoder_model_filename,
+        decoder_model_filename=args.decoder_model_filename,
+        joiner_model_filename=args.joiner_model_filename,
     )
 
-    decoder = ort.InferenceSession(
-        args.decoder_model_filename,
-        sess_options=session_opts,
-    )
-
-    joiner = ort.InferenceSession(
-        args.joiner_model_filename,
-        sess_options=session_opts,
-    )
-
-    joiner_encoder_proj = ort.InferenceSession(
-        args.joiner_encoder_proj_model_filename,
-        sess_options=session_opts,
-    )
-
-    joiner_decoder_proj = ort.InferenceSession(
-        args.joiner_decoder_proj_model_filename,
-        sess_options=session_opts,
-    )
-
-    sp = spm.SentencePieceProcessor()
-    sp.load(args.bpe_model)
-
     logging.info("Constructing Fbank computer")
     opts = kaldifeat.FbankOptions()
     opts.device = "cpu"
@@ -347,30 +386,27 @@ def main():
     )
 
     feature_lengths = torch.tensor(feature_lengths, dtype=torch.int64)
-
-    encoder_input_nodes = encoder.get_inputs()
-    encoder_out_nodes = encoder.get_outputs()
-    encoder_out, encoder_out_lens = encoder.run(
-        [encoder_out_nodes[0].name, encoder_out_nodes[1].name],
-        {
-            encoder_input_nodes[0].name: features.numpy(),
-            encoder_input_nodes[1].name: feature_lengths.numpy(),
-        },
-    )
+    encoder_out, encoder_out_lens = model.run_encoder(features, feature_lengths)
 
     hyps = greedy_search(
-        decoder=decoder,
-        joiner=joiner,
-        joiner_encoder_proj=joiner_encoder_proj,
-        joiner_decoder_proj=joiner_decoder_proj,
+        model=model,
         encoder_out=encoder_out,
         encoder_out_lens=encoder_out_lens,
-        context_size=args.context_size,
     )
     s = "\n"
+
+    symbol_table = k2.SymbolTable.from_file(args.tokens)
+
+    def token_ids_to_words(token_ids: List[int]) -> str:
+        text = ""
+        for i in token_ids:
+            text += symbol_table[i]
+        return text.replace("▁", " ").strip()
+
+    context_size = model.context_size
     for filename, hyp in zip(args.sound_files, hyps):
-        words = sp.decode(hyp)
-        s += f"{filename}:\n{words}\n\n"
+        words = token_ids_to_words(hyp[context_size:])
+        s += f"{filename}:\n{words}\n"
     logging.info(s)
 
     logging.info("Decoding Done")
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export-onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export-onnx.py
index a72472495..35d6b0556 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export-onnx.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/export-onnx.py
@@ -146,7 +146,7 @@ class OnnxEncoder(nn.Module):
         """
         Args:
           encoder:
-            A zipformer encoder.
+            A Zipformer encoder.
           encoder_proj:
             The projection layer for encoder from the joiner.
         """
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_check.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_check.py
index 72ad59a55..6c78ba70b 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_check.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_check.py
@@ -76,14 +76,8 @@ from zipformer import stack_states
 
 from icefall import is_module_available
 
-if not is_module_available("onnxruntime"):
-    raise ValueError("Please 'pip install onnxruntime' first.")
-
-import onnxruntime as ort
 import torch
 
-ort.set_default_logger_severity(3)
-
 
 def get_parser():
     parser = argparse.ArgumentParser(
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py
index 265684d18..715560c70 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py
@@ -333,7 +333,6 @@ class OnnxModel:
                 self.joiner.get_inputs()[1].name: decoder_out.numpy(),
             },
         )[0]
-        return torch.from_numpy(out)
 
         return torch.from_numpy(out)