Initial (unfinished) QKeras to QONNX converter

.gitignore

@@ -137,3 +137,6 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
+
+# IDE stuff
+.vscode

setup.cfg

@@ -53,6 +53,8 @@ install_requires =
     onnxruntime==1.11.1
     sigtools==2.0.3
     toposort>=1.5.0
+    tf2onnx==1.9.2
+    qkeras==0.9.0
 
 
 [options.packages.find]

src/qonnx/__init__.py

@@ -1,5 +1,7 @@
 import onnxruntime
 
+from qonnx import converters
+
 
 def reseed(newseed):
     onnxruntime.set_seed(newseed)

src/qonnx/converters/__init__.py

@@ -0,0 +1 @@
+from .keras import from_keras

src/qonnx/converters/keras.py

@@ -0,0 +1,148 @@
+import tensorflow as tf
+import tf2onnx
+from qkeras.utils import REGISTERED_LAYERS as QKERAS_LAYERS
+
+from finn.core.modelwrapper import ModelWrapper
+from qonnx.util.cleanup import cleanup_model
+
+from .qkeras.onnx import get_qkeras_onnx_handlers
+from .qkeras.qlayers import extract_quantizers_from_layer
+
+_unsupported_layers = [
+    # These require some extra work
+    "QBatchNormalization",
+    "QConv2DBatchnorm",
+    "QDepthwiseConv2DBatchnorm",
+]
+
+
+def _is_qkeras_model(model):
+    def iterate_model(model):
+        for layer in model.layers:
+            if isinstance(layer, tf.keras.Model):
+                found_qkeras = iterate_model(layer)
+                if found_qkeras:
+                    return True
+            elif layer.__class__.__name__ in QKERAS_LAYERS:
+                return True
+
+        return False
+
+    return iterate_model(model)
+
+
+def _check_supported_layers(model):
+    def iterate_model(model):
+        for layer in model.layers:
+            if isinstance(layer, tf.keras.Model):
+                iterate_model(layer)
+            elif layer.__class__.__name__ in _unsupported_layers:
+                raise Exception("Currently unsupported layer found in QKeras model: {}".format(layer.__class__.__name__))
+
+    iterate_model(model)
+
+
+def _strip_qkeras_model(model):
+    quantizers = {}
+
+    def extract_quantizers(layer):
+        keras_cls_name, layer_cfg, layer_quantizers = extract_quantizers_from_layer(layer)
+        if layer_quantizers:
+            layer_quantizers = {
+                k: None if v == "None" else v for k, v in layer_quantizers.items()
+            }  # Get rid of 'None' strings
+            quantizers[layer.name] = layer_quantizers
+
+        layer_class = tf.keras.layers.__dict__.get(keras_cls_name, None)
+        if layer_class is None:
+            raise Exception("Cannot create Keras layer from QKeras class {}".format(keras_cls_name))
+
+        return layer_class.from_config(layer_cfg)
+
+    stripped_model = tf.keras.models.clone_model(model, clone_function=extract_quantizers)
+
+    return stripped_model, quantizers
+
+
+def _convert_quantizers_to_nodes(onnx_model, quantizers_dict):
+
+    for node_name, quantizers in quantizers_dict.items():
+        print(node_name, quantizers)
+
+    for n in onnx_model.graph.node:
+        print(n)
+
+    return onnx_model.model
+
+
+def from_keras(
+    model,
+    input_signature=None,
+    opset=None,
+    custom_ops=None,
+    custom_op_handlers=None,
+    custom_rewriter=None,
+    inputs_as_nchw=None,
+    extra_opset=None,
+    shape_override=None,
+    target=None,
+    large_model=False,
+    output_path=None,
+):
+    """Convert a keras model to QONNX. The API follows the `from_keras` function of tf2onnx.
+
+    Args:
+        model: the tf.keras model we want to convert
+        input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input
+        opset: the opset to be used for the ONNX model, default is the latest
+        custom_ops: if a model contains ops not recognized by onnx runtime,
+            you can tag these ops with a custom op domain so that the
+            runtime can still open the model. Type is a dictionary `{op name: domain}`.
+        target: list of workarounds applied to help certain platforms
+        custom_op_handlers: dictionary of custom ops handlers
+        custom_rewriter: list of custom graph rewriters
+        extra_opset: list of extra opset's, for example the opset's used by custom ops
+        shape_override: dict with inputs that override the shapes given by tensorflow
+        inputs_as_nchw: transpose inputs in list from nchw to nhwc
+        large_model: use the ONNX external tensor storage format
+        output_path: save model to output_path
+
+    Returns:
+        An ONNX model_proto and an external_tensor_storage dict.
+    """
+
+    assert not large_model  # TODO for now, let's focus only on models that don't store tensors externally
+
+    if _is_qkeras_model(model):
+        _check_supported_layers(model)
+        keras_model, quantizers = _strip_qkeras_model(model)
+    else:
+        keras_model, quantizers = model, {}
+
+    qkeras_op_handlers = get_qkeras_onnx_handlers(quantizers)
+
+    if custom_op_handlers is not None:
+        qkeras_op_handlers.update(custom_op_handlers)
+
+    model_proto, external_storage = tf2onnx.convert.from_keras(
+        keras_model,
+        input_signature=input_signature,
+        opset=opset,
+        custom_ops=qkeras_op_handlers,
+        custom_op_handlers=qkeras_op_handlers,
+        custom_rewriter=custom_rewriter,
+        inputs_as_nchw=inputs_as_nchw,
+        extra_opset=extra_opset,
+        shape_override=shape_override,
+        target=target,
+        large_model=large_model,
+        output_path=None,
+    )
+
+    onnx_model = ModelWrapper(model_proto)
+    cleanup_model(onnx_model)
+
+    if output_path is not None:
+        onnx_model.save(output_path)
+
+    return onnx_model.model, external_storage

src/qonnx/converters/qkeras/onnx.py

@@ -0,0 +1,114 @@
+from tf2onnx.onnx_opset.math import DirectOp, MatMul
+from tf2onnx.onnx_opset.nn import BiasAdd, ConvOp
+
+from .quantizers import get_quant_params
+
+
+def get_qkeras_onnx_handlers(all_quantizers):
+    return {
+        "Conv2D": (conv2d_handler, ["Conv2D", all_quantizers]),
+        "MatMul": (dense_handler, ["MatMul", all_quantizers]),
+        "BiasAdd": (bias_handler, ["BiasAdd", all_quantizers]),
+        "Relu": (relu_handler, ["Relu", all_quantizers]),
+    }
+
+
+def _extract_node_name(onnx_name, keras_names):
+    for keras_name in keras_names:
+        match = "/" + keras_name + "/"
+        if match in onnx_name:
+            return keras_name
+
+    return None
+
+
+def qlayer_handler(ctx, node, name, args):
+    all_quantizers = args[0]
+    keras_name = _extract_node_name(name, all_quantizers.keys())
+    if not keras_name:
+        return  # Not found in quantizers, nothing to do
+    quantizers = all_quantizers[keras_name]
+
+    if quantizers.get("kernel_quantizer"):
+        weights = node.inputs[1].get_tensor_value(as_list=True)
+        kernel_quant_params = get_quant_params(weights, quantizers["kernel_quantizer"])
+        ctx.insert_new_node_on_input(
+            node, "Quant", node.input[1], name=node.name + "_kernel_quantizer", **kernel_quant_params, domain="qonnx"
+        )
+
+    if quantizers.get("bias_quantizer") and len(node.input) == 3:
+        bias = node.inputs[2].get_tensor_value(as_list=True)
+        bias_quant_params = get_quant_params(bias, quantizers["bias_quantizer"])
+        ctx.insert_new_node_on_input(
+            node, "Quant", node.input[2], name=node.name + "_bias_quantizer", **bias_quant_params, domain="qonnx"
+        )
+
+    if quantizers.get("activation"):
+        output_shapes = [ctx.get_shape(node.output[0])]
+        dtypes = [ctx.get_dtype(node.output[0])]
+        act_quant_params = get_quant_params(None, quantizers["activation"])
+        quant_act_node = ctx.make_node(
+            "Quant",
+            [node.output[0]],
+            shapes=output_shapes,
+            dtypes=dtypes,
+            name=node.name + "_activation_quantizer",
+            attr=act_quant_params,
+            domain="qonnx",
+        )
+        ctx.insert_node_on_output(quant_act_node, node.output[0])
+
+
+def qact_handler(ctx, node, name, args):
+    all_quantizers = args[0]
+    keras_name = _extract_node_name(name, all_quantizers.keys())
+    if not keras_name:
+        return  # Not found in quantizers, nothing to do
+    quantizers = all_quantizers[keras_name]
+
+    if quantizers.get("activation"):
+        output_shapes = [ctx.get_shape(node.output[0])]
+        dtypes = [ctx.get_dtype(node.output[0])]
+        act_quant_params = get_quant_params(None, quantizers["activation"])
+        quant_act_node = ctx.make_node(
+            "Quant",
+            [node.output[0]],
+            shapes=output_shapes,
+            dtypes=dtypes,
+            name=node.name + "_activation_quantizer",
+            attr=act_quant_params,
+            domain="qonnx",
+        )
+        ctx.insert_node_on_output(quant_act_node, node.output[0])
+
+
+def conv2d_handler(ctx, node, name, args):
+    ConvOp.any_version(11, ctx, node)
+    qlayer_handler(ctx, node, name, args)
+
+
+def dense_handler(ctx, node, name, args):
+    MatMul.version_1(ctx, node)
+    qlayer_handler(ctx, node, name, args)
+
+
+def bias_handler(ctx, node, name, args):
+    BiasAdd.version_1(ctx, node)
+
+    all_quantizers = args[0]
+    keras_name = _extract_node_name(name, all_quantizers.keys())
+    if not keras_name:
+        return  # Not found in quantizers, nothing to do
+    quantizers = all_quantizers[keras_name]
+
+    if quantizers.get("bias_quantizer"):
+        bias = node.inputs[1].get_tensor_value(as_list=True)
+        bias_quant_params = get_quant_params(bias, quantizers["bias_quantizer"])
+        ctx.insert_new_node_on_input(
+            node, "Quant", node.input[1], name=node.name + "_bias_quantizer", **bias_quant_params, domain="qonnx"
+        )
+
+
+def relu_handler(ctx, node, name, args):
+    DirectOp.version_1(ctx, node)
+    # qact_handler(ctx, node, name, args)