set_converter not working for tf2.X

[liob]

Hello,

I think set_converter is not working properly in tf.keras 2.2. Here is a minimal example:

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv3D, Layer
import keras2onnx
from keras2onnx import set_converter


class MyDenseLayer(Layer):
    ''' https://www.tensorflow.org/tutorials/customization/custom_layers '''
    def __init__(self, num_outputs):
        super(MyDenseLayer, self).__init__()
        self.num_outputs = num_outputs

    def build(self, input_shape):
        self.kernel = self.add_weight("kernel",
                                  shape=[int(input_shape[-1]),
                                  self.num_outputs])

    def call(self, input):
        return tf.matmul(input, self.kernel)
    
    def get_config(self):
        config = {'num_outputs':  self.num_outputs}
        base_config = super(MyDenseLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


# create model
input_tensor = Input(shape=(64, 64, 64, 1))
net = Conv3D(12, 3, activation='relu')(input_tensor)
net = MyDenseLayer(10)(net)
model = Model(input_tensor, net)


def convert_my_dense_layer(scope, operator, container):
    raise Exception('conversion function reached!')

set_converter(MyDenseLayer, convert_my_dense_layer)
onnx_model = keras2onnx.convert_keras(model, debug_mode=True)

I would expect, that the convert_my_dense_layer function is being called when encountering the MyDenseLayer instance during conversion. This is not the case.

Furthermore, the debug output is somewhat strange: If MyDenseLayer is not present in the model topology, than Conv3D is recognized as one unit:

Converting the operator (Identity): Identity Converting the operator (Identity1): Identity Converting the operator (Identity2): Identity Converting the operator (conv3d): <class 'tensorflow.python.keras.layers.convolutional.Conv3D'> Converting the operator (Identity3): Identity

If MyDenseLayer is present, than the log looks something like this:

Converting the operator (Identity): Identity Converting the operator (Identity1): Identity Converting the operator (Identity2): Identity Converting the operator (Identity3): Identity Converting the operator (model/my_dense_layer/MatMul): BatchMatMulV2 Converting the operator (model/conv3d/Relu): Relu Converting the operator (model/my_dense_layer/MatMul/ReadVariableOp): Identity Converting the operator (model/conv3d/BiasAdd): BiasAdd Converting the operator (model/my_dense_layer/MatMul/ReadVariableOp/resource): Const Converting the operator (model/conv3d/Conv3D): Conv3D Converting the operator (model/conv3d/BiasAdd/ReadVariableOp): Identity Converting the operator (model/conv3d/Conv3D/ReadVariableOp): Identity Converting the operator (model/conv3d/BiasAdd/ReadVariableOp/resource): Const Converting the operator (model/conv3d/Conv3D/ReadVariableOp/resource): Const Converting the operator (Identity4): Identity

[wenbingl]

Yes, it's a known issue that the tf 2.x subclassed layer doesn't support the custom conversion, only tensorflow op could be converted in the custom mode.

[liob]

@wenbingl I have backported the custom layer to tf 1.15. I see the same behavior as in tf 2.2 (the graph disintegrates into tensorflow ops). Does this mean, that set_converter is only working with keras.io?

If this is the case, is there any other way to implement custom layers?

[wenbingl]

with tf2.x, set_converter only supports to customize tf.op, with tf1.x, set_converter works on both tf.op and keras.layer. Is the above model source the issue happen?

[liob]

@wenbingl i did try to convert another dummy model. In this other model, I did use a tensorflow.keras.backend.zeros node. This did derail the conversion. After substituting the zeros node, the conversion function is being called. Sorry for the confusion.

[thoomi]

Do you have an idea what's necessary to support this? Changing the parser? Or are the keras layer converters not used at all for converting tf 2.2 models?

[liob]

@thoomi in the end, I was unable to achieve good results for my specific network topology and ended up writing my own custom onnx converter which works directly with tf.keras (2.x) layers: https://github.com/liob/naive_k2onnx

However, this is certainly not sustainable and I will switch to keras-onnx when there is an easy way to support custom tf.keras layers in 2.x.