Custom Masking value (#389)

* Adding broken test for using custom mask_value

* Allowing custom mask_value to be used

* Fixing dtypes in mask_value Equal operator

* Adding meaningful exception when the masking value is not specified

Co-authored-by: David Fan <[email protected]>

Author: cjermain

keras2onnx/_parser_1x.py

@@ -111,6 +111,9 @@ def on_parsing_keras_layer(graph, node_list, layer, kenode, model, varset, prefi
             mts_var = varset.get_local_variable_or_declare_one(mts_name, infer_variable_type(om_, varset.target_opset))
             operator.add_output_mask(mts_var)
 
+    if hasattr(layer, 'mask_value') and layer.mask_value is not None:
+        operator.mask_value = layer.mask_value
+
     cvt = get_converter(operator.type)
     if cvt is not None and hasattr(cvt, 'shape_infer'):
         operator.shape_infer = cvt.shape_infer

keras2onnx/common/intop.py

@@ -26,6 +26,7 @@ def __init__(self, onnx_name, scope, type, raw_operator, target_opset):
         self.input_masks = []
         self.outputs = []
         self.output_masks = []
+        self.mask_value = None
         self.nodelist = None
         self.is_evaluated = None
         self.target_opset = target_opset

keras2onnx/ke2onnx/main.py

@@ -68,8 +68,10 @@ def convert_keras_flatten(scope, operator, container):
 
 
 def _apply_not_equal(oopb, target_opset, operator):
+    if operator.mask_value is None:
+        raise ValueError("Masking value was not properly parsed for layer '{}'".format(operator.full_name))
     if target_opset >= 11:
-        equal_out = oopb.add_node('Equal', [operator.inputs[0].full_name, np.array([0], dtype='float32')],
+        equal_out = oopb.add_node('Equal', [operator.inputs[0].full_name, np.array([operator.mask_value], dtype='float32')],
                                   operator.full_name + 'mask')
         not_o = oopb.add_node('Not', equal_out,
                               name=operator.full_name + '_not')
@@ -78,7 +80,7 @@ def _apply_not_equal(oopb, target_opset, operator):
                              "the masking layer result may be incorrect if the model input is in range (0, 1.0).")
         equal_input_0 = oopb.add_node('Cast', [operator.inputs[0].full_name],
                                       operator.full_name + '_input_cast', to=6)
-        equal_out = oopb.add_node('Equal', [equal_input_0, np.array([0], dtype='int32')],
+        equal_out = oopb.add_node('Equal', [equal_input_0, np.array([operator.mask_value], dtype='int32')],
                                   operator.full_name + 'mask')
         not_o = oopb.add_node('Not', equal_out,
                               name=operator.full_name + '_not')

tests/test_layers.py

@@ -1788,6 +1788,21 @@ def test_masking(self):
         expected = model.predict(x)
         self.assertTrue(run_onnx_runtime(onnx_model.graph.name, onnx_model, x, expected, self.model_files))
 
+    @unittest.skipIf(is_tf2 and is_tf_keras, 'TODO')
+    def test_masking_value(self):
+        timesteps, features = (3, 5)
+        mask_value = 5.
+        model = Sequential([
+            keras.layers.Masking(mask_value=mask_value, input_shape=(timesteps, features)),
+            LSTM(8, return_state=False, return_sequences=False)
+        ])
+
+        onnx_model = keras2onnx.convert_keras(model, model.name)
+        x = np.random.uniform(100, 999, size=(2, 3, 5)).astype(np.float32)
+        x[1, :, :] = mask_value
+        expected = model.predict(x)
+        self.assertTrue(run_onnx_runtime(onnx_model.graph.name, onnx_model, x, expected, self.model_files))
+
     @unittest.skipIf(is_tf2 and is_tf_keras, 'TODO')
     def test_masking_custom(self):
         class MyPoolingMask(keras.layers.Layer):