how to export these onnx?

[lucasjinreal]

I tried export lstm-transducer, fails.

[csukuangfj]

It does not support LSTM with projections.

[lucasjinreal]

@csukuangfj How can I export these lstm models? I don't care it support or not, I just need a workable version

[csukuangfj]

Please see the doc

https://k2-fsa.github.io/icefall/recipes/librispeech/lstm_pruned_stateless_transducer.html#export-models

I posted it before in another thread.

[EmreOzkose]

In this version, models of pruned_transducer_stateless3 are used. In this script, models are exported separately, but projection parts of joiner should also be exported separately. You can use the below function to export joiner

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 model has two inputs:

        - encoder_out: a tensor of shape (N, encoder_out_dim)
        - decoder_out: a tensor of shape (N, decoder_out_dim)

    and has one output:

        - joiner_out: a tensor of shape (N, vocab_size)

    """
    encoder_out_dim = joiner_model.encoder_proj.weight.shape[1]
    decoder_out_dim = joiner_model.decoder_proj.weight.shape[1]
    encoder_out = torch.rand(1, 1, 1, encoder_out_dim, dtype=torch.float32)
    decoder_out = torch.rand(1, 1, 1, decoder_out_dim, dtype=torch.float32)

    project_input = False
    # Note: It uses torch.jit.trace() internally
    torch.onnx.export(
        joiner_model,
        (encoder_out, decoder_out, project_input),
        joiner_filename,
        verbose=False,
        opset_version=opset_version,
        input_names=["encoder_out", "decoder_out", "project_input"],
        output_names=["logit"],
        dynamic_axes={
            "encoder_out": {0: "N"},
            "decoder_out": {0: "N"},
            "logit": {0: "N"},
        },
    )
    torch.onnx.export(
        joiner_model.encoder_proj,
        (encoder_out.squeeze(0).squeeze(0)),
        str(joiner_filename).replace(".onnx", "_encoder_proj.onnx"),
        verbose=False,
        opset_version=opset_version,
        input_names=["encoder_out"],
        output_names=["encoder_proj"],
        dynamic_axes={
            "encoder_out": {0: "N"},
            "encoder_proj": {0: "N"},
        },
    )
    torch.onnx.export(
        joiner_model.decoder_proj,
        (decoder_out.squeeze(0).squeeze(0)),
        str(joiner_filename).replace(".onnx", "_decoder_proj.onnx"),
        verbose=False,
        opset_version=opset_version,
        input_names=["decoder_out"],
        output_names=["decoder_proj"],
        dynamic_axes={
            "decoder_out": {0: "N"},
            "decoder_proj": {0: "N"},
        },
    )
    logging.info(f"Saved to {joiner_filename}")

@csukuangfj , I might update onnx exporting script ?

[lucasjinreal]

@EmreOzkose So the model is conformer not lstm? Please update export script, I want have a tried on onnx.

[csukuangfj]

@csukuangfj , I might update onnx exporting script ?

Yes, please.

[lucasjinreal]

@EmreOzkose Hi, I wanna using wenet Chinese mode, how should I download pretrained model and convert toonnx? Does there any necessary to change thecode?

[EmreOzkose]

@EmreOzkose So the model is conformer not lstm? Please update export script, I want have a tried on onnx.

Yes, it is not LSTM. I made a PR.

@EmreOzkose Hi, I wanna using wenet Chinese mode, how should I download pretrained model and convert toonnx? Does there any necessary to change thecode?

Actually, I did experiments with only English pre-trained model. I have never worked on a Chinese model. If Chinese model doesn't have extra changes in greedy search, I think it will work.

[EmreOzkose]

You can export English models as below

git lfs install
git clone https://huggingface.co/csukuangfj/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13
cd exp
cp pretrained-iter-824000-avg-18.pt epoch-1.pt
cd path/to/icefall/egs/librispeech/ASR/
./pruned_transducer_stateless3/export.py\
 --exp-dir /path/to/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/exp \
 --bpe-model /path/to/icefall-asr-librispeech-pruned-transducer-stateless3-2022-05-13/data/lang_bpe_500/bpe.model \
 --epoch 1 \
 --avg 1 \
 --onnx 1

[csukuangfj]

This one is for non-streaming Conformer models. It should work for English as well as Chinese models.

[lucasjinreal]

@csukuangfj Hello, I got an error when load wenet model to inference:

 for decoder.embedding.weight: copying a param with shape torch.Size([5537, 512]) from checkpoint, the shape in current model is torch.Size([5539, 512]).

do u know why?

[lucasjinreal]

I get vocabsize from token got max + 1 : : vocab size: 5539

but model shape is 5537, why?

[csukuangfj]

Please show the complete code. It is hard to figure out what goes wrong without seeing the code.

[csukuangfj]

Please see
https://k2-fsa.github.io/icefall/model-export/index.html