crnn模型onnx转rknn报错

环境: ubuntu 18.04 / rknn-toolkit 1.3.2 / pytorch 1.2.0　

crnn模型是pytorch训练的，转成onnx模型，然后尝试转成rknn模型，在加载onnx模型时报错：
E Catch exception when loading onnx model: ../netCRNN_rknn.onnx!
E Traceback (most recent call last):
E   File "rknn/api/rknn_base.py", line 510, in rknn.api.rknn_base.RKNNBase.load_onnx
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 514, in rknn.base.RKNNlib.converter.convert_onnx.convert_onnx.__init__
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 176, in rknn.base.RKNNlib.converter.convert_onnx.onnx_shape_infer_engine.infer_shape
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 131, in rknn.base.RKNNlib.converter.convert_onnx.onnx_shape_infer_engine.__status_update
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 165, in rknn.base.RKNNlib.converter.convert_onnx.onnx_shape_infer_engine.__model_rebuild
E   File "rknn/base/RKNNlib/converter/tensor_model.py", line 330, in rknn.base.RKNNlib.converter.tensor_model.TensorModel.build_onnx_model
E KeyError: ''
Load model failed!

查询文档得知KeyError似乎是因为有op不支持，但是后面提示是空的，不知道是哪个op报错了，我把onnx的所有node都打印出来，统计了所有的op_type如下(onnx版本1.6.0)：
{'Concat', 'Relu', 'Reshape', 'Squeeze', 'MaxPool', 'Add', 'ConstantOfShape', 'Shape', 'Sub', 'Gather', 'Gemm', 'Conv', 'Transpose', 'LSTM', 'Mul', 'Div', 'Slice', 'Clip', 'BatchNormalization', 'Pad', 'Unsqueeze', 'Constant'}
不知道是哪个op不支持呢？

也尝试过直接pytorch模型转rknn,也是加载模型的时候报错了:

E Catch exception when loading pytorch model: ./netCRNN_39.pth!
E Traceback (most recent call last):
E   File "rknn/api/rknn_base.py", line 611, in rknn.api.rknn_base.RKNNBase.load_pytorch
E   File "rknn/base/RKNNlib/app/importer/import_pytorch.py", line 97, in rknn.base.RKNNlib.app.importer.import_pytorch.ImportPytorch.run
E   File "rknn/base/RKNNlib/converter/convert_pytorch.py", line 552, in rknn.base.RKNNlib.converter.convert_pytorch.convert_pytorch.__init__
E   File "/home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/jit/__init__.py", line 162, in load
E     cpp_module = torch._C.import_ir_module(cu, f, map_location, _extra_files)
E RuntimeError: [enforce fail at inline_container.cc:137] . PytorchStreamReader failed reading zip archive: failed finding central directory
E frame #0: c10::ThrowEnforceNotMet(char const*, int, char const*, std::string const&, void const*) + 0x47 (0x7f9b7f27ee17 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libc10.so)
E frame #1: caffe2::serialize:yTorchStreamReader::valid(char const*) + 0x6b (0x7f9b8220775b in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libtorch.so)
E frame #2: caffe2::serialize:yTorchStreamReader::init() + 0x9a (0x7f9b8220b20a in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libtorch.so)
E frame #3: caffe2::serialize:yTorchStreamReader:yTorchStreamReader(std::string const&) + 0x60 (0x7f9b8220e270 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libtorch.so)
E frame #4: torch::jit::import_ir_module(std::shared_ptr<torch::jit::script::CompilationUnit>, std::string const&, c10:ptional<c10:evice>, std::unordered_map<std::string, std::string, std::hash<std::string>, std::equal_to<std::string>, std::allocator<std::pair<std::string const, std::string> > >&) + 0x38 (0x7f9b832ed088 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libtorch.so)
E frame #5: <unknown function> + 0x4d69dc (0x7f9bc95ee9dc in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
E frame #6: <unknown function> + 0x1d3ef4 (0x7f9bc92ebef4 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
E frame #7: python() [0x50a635]
E <omitting python frames>
E frame #9: python() [0x507d64]
E frame #11: <unknown function> + 0x196b8 (0x7f9b3ce5c6b8 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/base/RKNNlib/converter/convert_pytorch.cpython-36m-x86_64-linux-gnu.so)
E frame #12: <unknown function> + 0x19ff7 (0x7f9b3ce5cff7 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/base/RKNNlib/converter/convert_pytorch.cpython-36m-x86_64-linux-gnu.so)
E frame #13: <unknown function> + 0x995d5 (0x7f9b3cedc5d5 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/base/RKNNlib/converter/convert_pytorch.cpython-36m-x86_64-linux-gnu.so)
E frame #14: python() [0x5946bc]
E frame #15: python() [0x54a005]
E frame #16: python() [0x5513d1]
E frame #17: <unknown function> + 0x6a8c (0x7f9b3e683a8c in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/base/RKNNlib/app/importer/import_pytorch.cpython-36m-x86_64-linux-gnu.so)
E frame #18: <unknown function> + 0xcd2e (0x7f9b3e689d2e in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/base/RKNNlib/app/importer/import_pytorch.cpython-36m-x86_64-linux-gnu.so)
E frame #19: <unknown function> + 0x1314c (0x7f9b3e69014c in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/base/RKNNlib/app/importer/import_pytorch.cpython-36m-x86_64-linux-gnu.so)
E frame #20: <unknown function> + 0x2041d (0x7f9be076b41d in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/api/rknn_base.cpython-36m-x86_64-linux-gnu.so)
E frame #21: <unknown function> + 0x2061e (0x7f9be076b61e in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/api/rknn_base.cpython-36m-x86_64-linux-gnu.so)
E frame #22: <unknown function> + 0x109436 (0x7f9be0854436 in /home/tan-nbk/files/env/rk3399pro/lib/python3.6/site-packages/rknn/api/rknn_base.cpython-36m-x86_64-linux-gnu.so)
E frame #24: python() [0x50a5c3]
E frame #26: python() [0x507d64]
E frame #27: python() [0x509a90]
E frame #28: python() [0x50a48d]
E frame #30: python() [0x507d64]
E frame #32: python() [0x634c82]
E frame #37: __libc_start_main + 0xe7 (0x7f9c03c6fb97 in /lib/x86_64-linux-gnu/libc.so.6)
Load pytorch model failed!

存onnx时，尝试把opset_version设置成9 试试。

onnx版本确认切换到1.4.1

leok 发表于 2020-8-6 16:00
onnx版本确认切换到1.4.1

pytorch模型导出的时候onnx opset_version设置为9, onnx版本切换为1.4.1,最后还是报一样的错误。
导出onnx模型的时候代码设置：

1. torch.onnx.export(net, input, 'netCRNN_rknn.onnx', export_params=True, verbose=True, opset_version=9)

复制代码
onnx转rknn代码:

1. # Create RKNN object
   
2.     rknn = RKNN(verbose=True)
   
3. 
   
4.     # pre-process config
   
5.     print('--> config model')
   
6.     rknn.config(channel_mean_value='0 0 0 1', reorder_channel='0 1 2')
   
7.     print('done')
   
8. 
   
9.     # Load onnx model
   
10.     print('--> Loading model')
    
11.     ret = rknn.load_onnx(model='../netCRNN_rknn.onnx')
    
12.     if ret != 0:
    
13.         print('Load model failed!')
    
14.         exit(ret)
    
15.     print('done')

复制代码
onnx模型所有node的op_type如下:
{'Gather', 'Conv', 'Clip', 'LSTM', 'Gemm', 'Constant', 'Pad', 'Transpose', 'Add', 'Reshape', 'Slice', 'ConstantOfShape', 'Shape', 'Concat', 'BatchNormalization', 'Relu', 'MaxPool', 'Mul', 'Squeeze', 'Sub', 'Unsqueeze', 'Div'}
报错:
E Catch exception when loading onnx model: ../netCRNN_rknn.onnx!
E Traceback (most recent call last):
E   File "rknn/api/rknn_base.py", line 510, in rknn.api.rknn_base.RKNNBase.load_onnx
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 514, in rknn.base.RKNNlib.converter.convert_onnx.convert_onnx.__init__
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 176, in rknn.base.RKNNlib.converter.convert_onnx.onnx_shape_infer_engine.infer_shape
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 131, in rknn.base.RKNNlib.converter.convert_onnx.onnx_shape_infer_engine.__status_update
E   File "rknn/base/RKNNlib/converter/convert_onnx.py", line 165, in rknn.base.RKNNlib.converter.convert_onnx.onnx_shape_infer_engine.__model_rebuild
E   File "rknn/base/RKNNlib/converter/tensor_model.py", line 330, in rknn.base.RKNNlib.converter.tensor_model.TensorModel.build_onnx_model
E KeyError: ''
Load model failed!