diff --git a/keras2onnx/_builtin.py b/keras2onnx/_builtin.py index 1cd10c78..0f12a22a 100644 --- a/keras2onnx/_builtin.py +++ b/keras2onnx/_builtin.py @@ -20,6 +20,7 @@ class TYPES: # tf-node types: Identity = 'Identity' Const = 'Const' + AddN = 'AddN' Any = 'Any' All = 'All' BatchMatMul = 'BatchMatMul' @@ -170,6 +171,15 @@ def convert_tf_identity(scope, operator, container): default_convert(scope, operator, container) +@converter_func(TYPES.AddN) +def convert_tf_addn(scope, operator, container): + oopb = OnnxOperatorBuilder(container, scope) + oopb.apply_op_with_output("apply_add", + operator.input_full_names, + operator.output_full_names, + name=operator.full_name + '_add') + + @converter_func(TYPES.BatchToSpaceND) def convert_tf_batch_to_space(scope, operator, container): node = operator.raw_operator diff --git a/tests/test_layers.py b/tests/test_layers.py index 3567ac92..6fc264b9 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -97,6 +97,21 @@ def test_keras_lambda(self): expected = model.predict(data) self.assertTrue(run_onnx_runtime('onnx_lambda', onnx_model, data, expected, self.model_files)) + def test_tf_addn(self): + input1 = Input(shape=(5, 3, 4), dtype=tf.float32) + input2 = Input(shape=(5, 3, 4), dtype=tf.float32) + sum = Lambda(tf.add_n)([input1, input2]) + model = keras.models.Model(inputs=[input1, input2], outputs=sum) + + onnx_model = keras2onnx.convert_keras(model, 'tf_add_n') + batch_data1_shape = (2, 5, 3, 4) + batch_data2_shape = (2, 5, 3, 4) + data1 = np.random.rand(*batch_data1_shape).astype(np.float32) + data2 = np.random.rand(*batch_data2_shape).astype(np.float32) + expected = model.predict([data1, data2]) + self.assertTrue( + run_onnx_runtime('tf_add_n', onnx_model, [data1, data2], expected, self.model_files)) + def test_tf_conv(self): model = Sequential() k = tf.constant(np.random.normal(loc=0.0, scale=1.0, size=(1, 2, 3, 5)).astype(np.float32))