Skip to content

Instantly share code, notes, and snippets.

@sloev
Created March 14, 2023 20:39
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save sloev/cae3e8d5327f737c19bc4465425b8c3b to your computer and use it in GitHub Desktop.
Save sloev/cae3e8d5327f737c19bc4465425b8c3b to your computer and use it in GitHub Desktop.
convert keras 2.3.1 model to onnx 1.13.1
import pickle
import json
import tensorflow as tf
from keras.models import load_model
from keras.preprocessing.sequence import pad_sequences
from keras.models import Sequential
from keras.layers import Dense, Embedding, LSTM
import tensorflow as tf
import keras
from tensorflow.keras.models import Model
import keras.backend as K
# keras==2.3.1 tensorflow==1.15.2 h5py==2.10.0 protobuf==3.6.1 tf2onnx
MODELS_FOLDER = "."
model = load_model(f"{MODELS_FOLDER}/model.hdf5")
K.set_learning_phase(0)
def keras_to_pb(model, output_filename, output_node_names):
"""
This is the function to convert the Keras model to pb.
Args:
model: The Keras model.
output_filename: The output .pb file name.
output_node_names: The output nodes of the network. If None, then
the function gets the last layer name as the output node.
"""
# Get the names of the input and output nodes.
in_name = model.layers[0].get_output_at(0).name.split(':')[0]
if output_node_names is None:
output_node_names = [model.layers[-1].get_output_at(0).name.split(':')[0]]
sess = keras.backend.get_session()
# The TensorFlow freeze_graph expects a comma-separated string of output node names.
output_node_names_tf = ','.join(output_node_names)
frozen_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
sess.graph_def,
output_node_names)
sess.close()
wkdir = ''
tf.train.write_graph(frozen_graph_def, wkdir, output_filename, as_text=False)
return in_name, output_node_names
# Convert the Keras ResNet-50 model to a .pb file
in_tensor_name, out_tensor_names = keras_to_pb(model, "model.pb", None)
print("in", in_tensor_name)
print("out",out_tensor_names)

in python 3.7.9:

pip install keras==2.3.1   tensorflow==1.15.2  h5py==2.10.0      protobuf==3.6.1

run keras_to_pb.py and remember output names then in a python 3.10.x shell run:

pip install tf2onnx
python -m tf2onnx.convert  --input model.pb --inputs Tokenized_Sent_Input:0 --outputs Output_Predictions/Softmax:0 --output model_from_pb.onnx
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment