Here I show how to convert Keras models to the ONNX format using the kera2onnx package.
In this article in our series about using portable neural networks in 2020, you’ll learn how to convert a Keras model to the portable ONNX format.
Since ONNX is not a framework for building and training models, I will start with a brief introduction to Keras. This will be useful for engineers that are starting from scratch and are considering Keras as a framework for building and training models.
A Brief Introduction to Keras
Keras was designed to be an interface for building neural networks. This means it does not contain a runtime for training and serving models. Its interface was designed with human users in mind — it is high level and intuitive to use. When you look at the code that creates a model in Keras, it is easy to see all the layers involved and what they do.
Keras started out as a research project written by a Google engineer. Consequently, it eventually found its way into TensorFlow, so if you have 2.0 installed then you already have Keras.
A Quick Look at a Model
The code below creates and trains a model that predicts the digits in the MNIST dataset. (It is an excerpt from the full end-to-end demo that accompanies this article series.) In the next section, we will convert this model to the ONNX format.
def build_model():
model = keras.Sequential(
[
keras.Input(shape=input_shape),
layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Flatten(),
layers.Dropout(0.5),
layers.Dense(num_classes, activation="softmax"),
]
)
return model
def train_model(model, x_train, y_train):
model.compile(loss="categorical_crossentropy",
optimizer="adam",
metrics=["accuracy"])
model.fit(x_train, y_train,
batch_size=batch_size, epochs=epochs,
validation_split=0.1)
Installing and Importing the Converter
Before converting your Keras models to ONNX, you will need to install the keras2onnx package as it is not included with either Keras or TensorFlow. The following command installs the Keras to ONNX conversion utility:
pip install keras2onnx
Once installed, the converter can be imported into your modules using the following import:
import keras2onnx
Converting Keras Models to ONNX
Converting your Keras model to ONNX is as simple as running the function shown below. The only lines of code that are mandatory are the lines that convert the model and save the converted model to the file system. What is noteworthy about the keras2onnx converter is that it only requires the model as a parameter. This makes conversions easy and less error-prone. Other converters for other frameworks require information about the input to the model which is easy to get wrong.
def export_to_onnx(model):
onnx_model = keras2onnx.convert_keras(model, model.name)
meta = onnx_model.metadata_props.add()
meta.key = "creation_date"
meta.value = datetime.datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
meta = onnx_model.metadata_props.add()
meta.key = "author"
meta.value = 'keithpij'
onnx_model.doc_string = 'MNIST model'
onnx_model.model_version = 3
keras2onnx.save_model(onnx_model, ONNX_MODEL_FILE)
The code that adds metadata to the model is a best practice. As the data you use to train your model evolves, so will your model. Therefore, it is a good idea to add metadata to your model so that you can distinguish it from previous models. The example above adds a brief description of the model to the doc_string
property and sets the version. creation_date
and author
are custom properties added to the metadata_props
property bag. You are free to create as many custom properties using this property bag. Unfortunately, the model_version
property requires an integer or long so you will not be able to version it like your services using major.minor.revision syntax.
Summary and Next Steps
The keras2onnx package is true to the design goals of Keras itself. It is intuitive and easy to use. In this article, I provided a brief overview of Keras for those looking for a deep learning framework for building and training neural networks. I then showed how to convert Keras models to the ONNX format using the kera2onnx package.
Since the purpose of this article was to demonstrate converting Keras models to the ONNX format, I did not go into detail building and training Keras models. The code sample for this post contains code that explores Keras itself. The keras_mnist.py module is a full end-to-end demo that shows how to load the data, explore the images, and train the model.
Next, we’ll look at converting a PyTorch model to ONNX.
References
Keith is a sojourner in the software industry. He has over 30 years of experience building and bringing applications to market. He has worked for startups and large enterprises in roles ranging from tech lead to business development manager. He is currently a senior engineer on BNY Mellon's Distribution Analytics team where he is building data pipelines from on-premise data sources to the cloud.