问题场景
pytorch->onnx->caffe2出现的问题
在onnx->caffe2时会出现
Traceback (most recent call last):
File "/usr/local/bin/convert-onnx-to-caffe2", line 8, in <module>
sys.exit(onnx_to_caffe2())
File "/usr/local/lib/python3.6/dist-packages/click/core.py", line 764, in __call__
return self.main(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/click/core.py", line 717, in main
rv = self.invoke(ctx)
File "/usr/local/lib/python3.6/dist-packages/click/core.py", line 956, in invoke
return ctx.invoke(self.callback, **ctx.params)
File "/usr/local/lib/python3.6/dist-packages/click/core.py", line 555, in invoke
return callback(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/caffe2/python/onnx/bin/conversion.py", line 87, in onnx_to_caffe2
init_net, predict_net = c2.onnx_graph_to_caffe2_net(onnx_model_proto)
File "/usr/local/lib/python3.6/dist-packages/caffe2/python/onnx/backend.py", line 921, in onnx_graph_to_caffe2_net
return cls._onnx_model_to_caffe2_net(model, device=device, opset_version=opset_version, include_initializers=True)
File "/usr/local/lib/python3.6/dist-packages/caffe2/python/onnx/backend.py", line 876, in _onnx_model_to_caffe2_net
onnx_model = onnx.utils.polish_model(onnx_model)
File "/usr/local/lib/python3.6/dist-packages/onnx/utils.py", line 21, in polish_model
model = onnx.optimizer.optimize(model)
File "/usr/local/lib/python3.6/dist-packages/onnx/optimizer.py", line 55, in optimize
optimized_model_str = C.optimize(model_str, passes)
IndexError: _Map_base::at
修复
在pytorch导出的时候设置参数
torch.onnx._export(model, x, "model.onnx", export_params=True, keep_initializers_as_inputs=True)
原因
原因是导出模型时使用了CPU,导入caffe2时默认使用CUDA,导致错误。
参考地址
https://github.com/onnx/onnx/issues/1385