如何将使用Keras模型训练的Tensorflow 2. *转换为.onnx格式?

时间:2020-04-10 12:03:13

标签: python opencv tensorflow keras onnx

我将Python 3.7.4与TensorFlow 2.0和Keras 2.2.4-tf结合使用来训练自己的CNN模型。一切顺利。我可以使用例如model.save(my_model),然后在其他Python脚本中使用它。当我想在OpenCV中使用经过训练的模型及其C ++中的DNN模块时出现问题。 cv :: dnn:readNetFromTensorflow(model.pb,model.pbtxt),您会看到两个参数,而我无法获取第二个.pbtxt文件。因此,由于其灵活性,我决定使用.onnx格式。问题在于现有库keras2onnx仅采用TensorFlow 1. *中的模型,我想避免使用它。转换它的代码示例如下:

import tensorflow as tf
import onnx
import keras2onnx
model = tf.keras.models.load_model(my_model_folder_path)
onnx_model = keras2onnx.convert_keras(model, model.name)
onnx.save_model(onnx_model, model_name_onnx)

还有其他方法可以将这种模型转换为onnx格式吗?

2 个答案:

答案 0 :(得分:4)

最新版本的keras2onnx(在github主版本中)支持TensorFlow2。

您可以这样安装它:

pip install git+https://github.com/microsoft/onnxconverter-common
pip install git+https://github.com/onnx/keras-onnx

答案 1 :(得分:1)

您需要创建一个可以容纳ONNX对象的文件。访问https://github.com/onnx/tutorials/blob/master/tutorials/OnnxTensorflowExport.ipynb

import tensorflow as tf
import onnx
import keras2onnx
model = tf.keras.models.load_model('Model.h5')
onnx_model = keras2onnx.convert_keras(model, model.name)

file = open("Sample_model.onnx", "wb")
file.write(onnx_model.SerializeToString())
file.close()
相关问题