我在python中有一个基本的tensorflow模型,我想将其转换为onnx文件

时间:2019-12-11 16:58:46

标签: python tensorflow tensorflow2.0 onnx

我在python tensorflow中有基本模型,我想将其保存到onnx文件中,我该怎么做。我尝试使用onnx.save函数时遇到错误。

File "tenserflowbase.py", line 21, in <module> onnx.save(trained_model,'model.onxx') 
File "C:\Users\Parag_IK\Anaconda3\lib\site-packages\onnx\__init__.py", line 184, in save_model proto = write_external_data_tensors(proto, basepath)                                                                                                                                                                                                                                                           
File "C:\Users\Parag_IK\Anaconda3\lib\site packages\onnx\external_data_helper.py", line 225, in write_external_data_tensors                                               
for tensor in _get_all_tensors(model):                                                                                                                              
File "C:\Users\Parag_IK\Anaconda3\lib\site packages\onnx\external_data_helper.py", line 170, in _get_initializer_tensors                                                 
for initializer in onnx_model_proto.graph.initializer:                                                                                                              
AttributeError: 'History' object has no attribute 'graph'** 

我的代码如下:

import onnx
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

tf.logging.set_verbosity(tf.logging.ERROR)   

mar_budget = np.array([60, 80,  100  , 30, 50, 20, 90,  10],  dtype=float)
subs_gained = np.array([160, 200, 240, 100, 140, 80, 220, 60],  dtype=float)

for i, c in enumerate(mar_budget):
  print("{} Market budget = {} new subscribers gained".format(c, subs_gained[i]))


X_train, X_test, y_train, y_test = train_test_split(mar_budget, subs_gained, 
   random_state=42, train_size=0.8, test_size=0.2)

layer_0 = tf.keras.layers.Dense(units=1, input_shape=[1])
model = tf.keras.Sequential([layer_0])
model = tf.keras.Sequential([layer_0])

model.compile(loss='mean_squared_error', optimizer=tf.keras.optimizers.Adam(0.1))

trained_model = model.fit(X_train, y_train, epochs=1000, verbose=False)      

onnx.save(trained_model,'model.onxx')

print("Finished training the model")   
print(model.predict([80.0]))

1 个答案:

答案 0 :(得分:1)

我认为,如果我没记错的话,为了使用onnx.save(),模型应该在onnx函数中创建图形。

所以我建议使用tf2onnx库,该库具有将tf会话图转换为onnx图的功能。

onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph, ...) 

例如,完整代码为:

import tensorflow as tf
import tf2onnx

with tf.Session() as sess:
    x = tf.placeholder(tf.float32, [2, 3], name="input")
    x_ = tf.add(x, x)
    _ = tf.identity(x_, name="output")
    onnx_graph = tf2onnx.tfonnx.process_tf_graph(sess.graph, input_names=["input:0"], output_names=["output:0"])
    model_proto = onnx_graph.make_model("test")
    with open("/tmp/model.onnx", "wb") as f:
        f.write(model_proto.SerializeToString())

希望这会有所帮助。

参考:tensorflow-onnx