tf.keras to onnx

to_onnx.py

import tf2onnx
from tensorflow.python.keras.saving import saving_utils as _saving_utils

def to_onnx(model, output=None):
    function = _saving_utils.trace_model_call(model)
    concrete_func = function.get_concrete_function()

    input_names = [input_tensor.name for input_tensor in concrete_func.inputs
                   if input_tensor.dtype != tf.dtypes.resource]
    output_names = [output_tensor.name for output_tensor in concrete_func.outputs
                    if output_tensor.dtype != tf.dtypes.resource]

    frozen_graph = tf2onnx.tf_loader.from_function(concrete_func, input_names, output_names)
    with tf.Graph().as_default() as tf_graph:
        tf.import_graph_def(frozen_graph, name='')

    g = tf2onnx.tfonnx.process_tf_graph(tf_graph, opset=11, input_names=input_names, output_names=output_names)
    onnx_graph = tf2onnx.optimizer.optimize_graph(g)
    model_proto = onnx_graph.make_model("test model")
    if output:
        tf2onnx.utils.save_protobuf(output, model_proto)

    inputs = [n.name for n in model_proto.graph.input]
    outputs = [n.name for n in model_proto.graph.output]

    return model_proto, inputs, outputs