Making Keras Models Portable Using ONNX

• Download source - 547.1 KB

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):
   # convert to onnx model
   onnx_model = keras2onnx.convert_keras(model, model.name)

   # Add metadata to the ONNX model.
   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  # This must be an integer or long.
   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

• https://www.tensorflow.org/
• https://keras.io/
• https://github.com/onnx/keras-onnx
• https://github.com/keithpij/onnx-lab