ONNX Runtime for Keras

The following demonstrates how to compute the predictions of a pretrained deep learning model obtained from keras with onnxruntime. The conversion requires keras, tensorflow, keras-onnx, onnxmltools but then only onnxruntime is required to compute the predictions.

import os
if not os.path.exists('dense121.onnx'):
    from keras.applications.densenet import DenseNet121
    model = DenseNet121(include_top=True, weights='imagenet')

    from keras2onnx import convert_keras
    onx = convert_keras(model, 'dense121.onnx')
    onx.ir_version = 6
    with open("dense121.onnx", "wb") as f:
        f.write(onx.SerializeToString())

Let’s load an image (source: wikipedia).

from keras.preprocessing.image import array_to_img, img_to_array, load_img
img = load_img('Sannosawa1.jpg')
ximg = img_to_array(img)

import matplotlib.pyplot as plt
plt.imshow(ximg / 255)
plt.axis('off')

Out:

(-0.5, 1279.5, 959.5, -0.5)

Let’s load the model with onnxruntime.

import onnxruntime as rt
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidGraph

try:
    sess = rt.InferenceSession('dense121.onnx')
    ok = True
except (InvalidGraph, TypeError, RuntimeError) as e:
    # Probably a mismatch between onnxruntime and onnx version.
    print(e)
    ok = False

if ok:
    print("The model expects input shape:", sess.get_inputs()[0].shape)
    print("image shape:", ximg.shape)

Out:

The model expects input shape: ['N', 224, 224, 3]
image shape: (960, 1280, 3)

Let’s resize the image.

if ok:
    from skimage.transform import resize
    import numpy

    ximg224 = resize(ximg / 255, (224, 224, 3), anti_aliasing=True)
    ximg = ximg224[numpy.newaxis, :, :, :]
    ximg = ximg.astype(numpy.float32)

    print("new shape:", ximg.shape)

Out:

new shape: (1, 224, 224, 3)

Let’s compute the output.

if ok:
    input_name = sess.get_inputs()[0].name
    res = sess.run(None, {input_name: ximg})
    prob = res[0]
    print(prob.ravel()[:10])  # Too big to be displayed.

Out:

[2.0847994e-05 9.0344076e-07 1.6248664e-06 4.8085840e-06 6.5068948e-06
 9.4026956e-07 1.9977851e-06 4.6639366e-07 9.4333046e-07 3.2267365e-06]

Let’s get more comprehensive results.

if ok:
    from keras.applications.densenet import decode_predictions
    decoded = decode_predictions(prob)

    import pandas
    df = pandas.DataFrame(decoded[0], columns=["class_id", "name", "P"])
    print(df)

Out:

Downloading data from https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json

 8192/35363 [=====>........................] - ETA: 0s
40960/35363 [==================================] - 0s 0us/step
    class_id           name         P
0  n09468604         valley  0.673279
1  n09193705            alp  0.267428
2  n09399592     promontory  0.013859
3  n09246464          cliff  0.013251
4  n03792972  mountain_tent  0.007756