wenet-e2e / wenet

Production First and Production Ready End-to-End Speech Recognition Toolkit
https://wenet-e2e.github.io/wenet/
Apache License 2.0
4.07k stars 1.06k forks source link

streaming onnx export #2586

Open chopperWang opened 1 month ago

chopperWang commented 1 month ago

当我想导出自己训练的u2pp_conformer的流式模型时报错,请问这是什么原因呢?

Traceback (most recent call last):
  File "/mnt/workspace/wenet/examples/aishell/s0/wenet/bin/export_onnx_gpu.py", line 1228, in <module>
    onnx_config = export_enc_func(model, configs, args, logger,
  File "/mnt/workspace/wenet/examples/aishell/s0/wenet/bin/export_onnx_gpu.py", line 943, in export_online_encoder
    torch.onnx.export(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 516, in export
    _export(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1596, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1135, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1011, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 915, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py", line 1285, in _get_trace_graph
    outs = ONNXTracedModule(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py", line 133, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_trace.py", line 124, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1508, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/mnt/workspace/wenet/examples/aishell/s0/wenet/bin/export_onnx_gpu.py", line 172, in forward
    xs, _, new_att_cache, new_cnn_cache = layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1508, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/mnt/workspace/wenet/wenet/transformer/encoder_layer.py", line 225, in forward
    x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1508, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/mnt/workspace/wenet/wenet/transformer/attention.py", line 380, in forward
    k, v, new_cache = self._update_kv_and_cache(k, v, cache)
  File "/mnt/workspace/wenet/wenet/transformer/attention.py", line 196, in _update_kv_and_cache
    key_cache, value_cache = cache
ValueError: too many values to unpack (expected 2)
chopperWang commented 1 month ago

似乎是v3.1.0才有的问题,版本切回到v3.0.1就可以导出了