Support optional operators (onnx#980)

* add optional data type support

Signed-off-by: degaochu <[email protected]>

* fix seq test model issue

Signed-off-by: degaochu <[email protected]>

* update support status

Signed-off-by: degaochu <[email protected]>

* support identity opset 16

Signed-off-by: degaochu <[email protected]>

* fix issue

Signed-off-by: degaochu <[email protected]>

* fix issues for onnx 1.9.0

Signed-off-by: degaochu <[email protected]>

* fix attr proto cross version issue

Signed-off-by: degaochu <[email protected]>

* update opset list and status doc

Signed-off-by: degaochu <[email protected]>

Author: chudegao

doc/support_status.md

@@ -240,7 +240,10 @@ def run_node(cls, node, inputs, device='CPU', outputs_info=None, **kwargs):
     node = OnnxNode(node)
     input_tensors = []
     for i in inputs:
-      input_tensors.append(tf.constant(i))
+      if i is None:
+        input_tensors.append(i)
+      else:
+        input_tensors.append(tf.constant(i))
 
     if isinstance(inputs, dict):
       feed_dict_raw = inputs
@@ -252,7 +255,15 @@ def run_node(cls, node, inputs, device='CPU', outputs_info=None, **kwargs):
     input_dict = {}
     for k, v in feed_dict_raw.items():
       if isinstance(v, list):
-        input_dict[k] = [tf.constant(x) for x in v]
+        list_input = []
+        for x in v:
+          if x is None:
+            list_input.append(x)
+          else:
+            list_input.append(tf.constant(x))
+        input_dict[k] = list_input
+      elif v is None:  # keep None for empty optional data
+        input_dict[k] = v
       else:
         input_dict[k] = tf.constant(v)
 

onnx_tf/backend.py

@@ -78,6 +78,7 @@ def run(self, inputs, **kwargs):
     if isinstance(inputs, dict):
       feed_dict = inputs
     elif isinstance(inputs, list) or isinstance(inputs, tuple):
+      #to do: handle input is seq(optional) and optional element is empty(if/loop opset 16)
       if len(self.inputs) != len(inputs):
         raise RuntimeError('Expected {} values for uninitialized '
                            'graph inputs ({}), but got {}.'.format(
@@ -91,7 +92,15 @@ def run(self, inputs, **kwargs):
     input_dict = {}
     for k, v in feed_dict.items():
       if isinstance(v, list):
-        input_dict[k] = [tf.constant(x) for x in v]
+        list_input = []
+        for x in v:
+          if x is None:
+            list_input.append(x)
+          else:
+            list_input.append(tf.constant(x))
+        input_dict[k] = list_input
+      elif v is None:  # keep None for empty optional data
+        input_dict[k] = v
       else:
         input_dict[k] = tf.constant(v)
 
@@ -109,6 +118,8 @@ def run(self, inputs, **kwargs):
         o_values.append(v_list)
       elif isinstance(output_values[o_name], tf.Tensor):
         o_values.append(output_values[o_name].numpy())
+      elif isinstance(output_values[o_name], tf.RaggedTensor):
+        o_values.append([i for i in output_values[o_name].numpy()])
       else:
         o_values.append(output_values[o_name])
 

onnx_tf/backend_rep.py

@@ -1,4 +1,5 @@
 from onnx_tf.common import IS_PYTHON3
+from onnx_tf.common.legacy import legacy_opset_pre_ver
 
 
 def convert_tf(attr):
@@ -82,5 +83,7 @@ def __convert_onnx_attribute_proto(attr_proto):
     return str_list
   elif attr_proto.HasField('sparse_tensor'):
     return attr_proto.sparse_tensor
+  elif not legacy_opset_pre_ver(15) and attr_proto.HasField('tp'):
+    return attr_proto.tp
   else:
     raise ValueError("Unsupported ONNX attribute: {}".format(attr_proto))

onnx_tf/common/attr_converter.py

@@ -18,9 +18,17 @@ def version_13(cls, node, **kwargs):
     return [cls.make_tensor_from_onnx_node(node, **kwargs)]
 
   @classmethod
-  def version_14(cls, node, **kwargs):
+  def _common(cls, node, **kwargs):
     x = kwargs["tensor_dict"][node.inputs[0]]
     if isinstance(x, (list, tuple)):
       return [tf.identity_n(x)]
     else:
       return [tf.identity(x)]
+
+  @classmethod
+  def version_14(cls, node, **kwargs):
+    return cls._common(node, **kwargs)
+
+  @classmethod
+  def version_16(cls, node, **kwargs):
+    return cls._common(node, **kwargs)

onnx_tf/handlers/backend/identity.py

@@ -0,0 +1,13 @@
+from onnx_tf.handlers.backend_handler import BackendHandler
+from onnx_tf.handlers.handler import onnx_op
+
+
+@onnx_op("Optional")
+class Optional(BackendHandler):
+
+  @classmethod
+  def version_15(cls, node, **kwargs):
+    if len(node.inputs) > 0:
+        return [kwargs["tensor_dict"][node.inputs[0]]]
+    else:
+        return [None]

onnx_tf/handlers/backend/optional.py

@@ -0,0 +1,13 @@
+from onnx_tf.handlers.backend_handler import BackendHandler
+from onnx_tf.handlers.handler import onnx_op
+
+
+@onnx_op("OptionalGetElement")
+class OptionalGetElement(BackendHandler):
+
+  @classmethod
+  def version_15(cls, node, **kwargs):
+    if len(node.inputs) > 0:
+        return [kwargs["tensor_dict"][node.inputs[0]]]
+    else:
+        raise RuntimeError("No element value!.")

onnx_tf/handlers/backend/optional_get_element.py

@@ -0,0 +1,13 @@
+from onnx_tf.handlers.backend_handler import BackendHandler
+from onnx_tf.handlers.handler import onnx_op
+
+
+@onnx_op("OptionalHasElement")
+class OptionalHasElement(BackendHandler):
+
+  @classmethod
+  def version_15(cls, node, **kwargs):
+    if len(node.inputs) > 0 and kwargs["tensor_dict"][node.inputs[0]] is not None:
+        return [True]
+    else:
+        return [False]

onnx_tf/handlers/backend/optional_has_element.py

@@ -68,7 +68,7 @@
     'HardSigmoid': [1, 6],
     'HardSwish': [14],
     'Hardmax': [1, 11, 13],
-    'Identity': [1, 13, 14],
+    'Identity': [1, 13, 14, 16],
     'If': [1, 11, 13],
     'ImageScaler': [1],
     'Imputer': [],
@@ -109,9 +109,9 @@
     'Not': [1],
     'OneHot': [9, 11],
     'OneHotEncoder': [],
-    'Optional': [],
-    'OptionalGetElement': [],
-    'OptionalHasElement': [],
+    'Optional': [15],
+    'OptionalGetElement': [15],
+    'OptionalHasElement': [15],
     'Or': [1, 7],
     'PRelu': [1, 6, 7, 9],
     'Pad': [1, 2, 11, 13],

onnx_tf/opset_version.py

@@ -364,16 +364,16 @@ def test_if_with_sequence(self):
         'c': c,
         'cond': np.array(True, dtype=np.bool)
     })
-    np.testing.assert_almost_equal(output['S_final'].values[:2], a)
-    np.testing.assert_almost_equal(output['S_final'].values[2:], b)
+    np.testing.assert_almost_equal(output['S_final'][0], a)
+    np.testing.assert_almost_equal(output['S_final'][1], b)
     output = tf_rep.run({
         'a': a,
         'b': b,
         'c': c,
         'cond': np.array(False, dtype=np.bool)
     })
-    np.testing.assert_almost_equal(output['S_final'].values[:2], a)
-    np.testing.assert_almost_equal(output['S_final'].values[2:], c)
+    np.testing.assert_almost_equal(output['S_final'][0], a)
+    np.testing.assert_almost_equal(output['S_final'][1], c)
 
   def test_initializer(self):
     if legacy_onnx_pre_ver(1, 2):
@@ -460,10 +460,10 @@ def test_loop_with_sequence(self):
                                   outputs=[s_final_out])
     tf_rep = prepare(helper.make_model(graph_def))
     output = tf_rep.run({'a': a, 'b': b, 'M': M, 'cond_init': cond})
-    np.testing.assert_almost_equal(output['S_final'].values[:2], a)
-    np.testing.assert_almost_equal(output['S_final'].values[2:3], b)
-    np.testing.assert_almost_equal(output['S_final'].values[3:4], b)
-    np.testing.assert_almost_equal(output['S_final'].values[4:5], b)
+    np.testing.assert_almost_equal(output['S_final'][0], a)
+    np.testing.assert_almost_equal(output['S_final'][1], b)
+    np.testing.assert_almost_equal(output['S_final'][2], b)
+    np.testing.assert_almost_equal(output['S_final'][3], b)
 
   def test_pow_bfloat16(self):
     X1 = np.array([1, 2, 3]).astype(np.float32)

test/backend/test_model.py

@@ -4198,7 +4198,28 @@ def test_where(self):
     output = run_node(node_def, [c, x, y])
     np.testing.assert_almost_equal(output["Z"], np.where(c, x, y))
 
+  def test_optional(self):
+    if legacy_opset_pre_ver(16):
+      raise unittest.SkipTest("ONNX version {} doesn't support Optional.".format(
+          defs.onnx_opset_version()))
+    ten_in_tp = helper.make_tensor_type_proto(TensorProto.FLOAT, shape=[5])
+    opt_in_tp = helper.make_optional_type_proto(ten_in_tp)
+    node_def = helper.make_node("Optional", ["X"], ["Y"], type=opt_in_tp)
+    x = self._get_rnd_float32(-3.0, 3.0, [5])
+    output = run_node(node_def, [x])
+    np.testing.assert_almost_equal(output["Y"], x)
 
+  def test_optional_empty(self):
+    if legacy_opset_pre_ver(16):
+      raise unittest.SkipTest("ONNX version {} doesn't support Optional.".format(
+          defs.onnx_opset_version()))
+    ten_in_tp = helper.make_tensor_type_proto(TensorProto.FLOAT, shape=[5])
+    opt_in_tp = helper.make_optional_type_proto(ten_in_tp)
+    node_def = helper.make_node("Optional", ["X"], ["Y"], type=opt_in_tp)
+    x = None
+    output = run_node(node_def, [x])
+    np.testing.assert_equal(output["Y"], x)
+    
 if __name__ == '__main__':
   unittest.main()
 

test/backend/test_node.py

@@ -76,9 +76,6 @@ def get_onnx_supported_ops():
 if legacy_opset_pre_ver(13):
   backend_test.exclude(r'test_cumsum_[a-z,_]*')
 
-# Currently ONNX's backend test runner does not support sequence as input/output
-backend_test.exclude(r'test_if_seq[a-z,_]*')
-
 # TF session run does not support sequence/RaggedTensor as model inputs
 backend_test.exclude(r'test_loop13_seq[a-z,_]*')
 
@@ -149,9 +146,10 @@ def get_onnx_supported_ops():
 backend_test.exclude(r'test_bernoulli_double_expanded[a-z,_]*')
 backend_test.exclude(r'test_bernoulli_seed_expanded[a-z,_]*')
 
-# Exclude optional_get_element, test_optional_has_element tests
-backend_test.exclude(r'test_optional_get_element[a-z,_]*')
-backend_test.exclude(r'test_optional_has_element[a-z,_]*')
+# # onnx backend test support seq from 1.11 #3731
+if legacy_opset_pre_ver(16):
+  backend_test.exclude(r'test_optional_get_element[a-z,_]*')
+  backend_test.exclude(r'test_optional_has_element[a-z,_]*')
 
 # Exclude BatchNormalization with training_mode=1 tests
 backend_test.exclude(r'test_batchnorm_epsilon_training_mode[a-z,_]*')
@@ -161,8 +159,13 @@ def get_onnx_supported_ops():
 if legacy_opset_pre_ver(15):
   backend_test.exclude(r'[a-z,_]*identity_sequence_[a-z,_]*')
 
-# Exclude tests with optional inputs
-backend_test.exclude(r'[a-z,_]*identity_opt_[a-z,_]*')
+# onnx backend test support seq from 1.11 #3731
+if legacy_opset_pre_ver(16):
+  backend_test.exclude(r'[a-z,_]*identity_opt_[a-z,_]*')
+
+if legacy_opset_pre_ver(16):
+  backend_test.exclude(r'test_if_seq[a-z,_]*')
+
 backend_test.exclude(r'[a-z,_]*if_opt_[a-z,_]*')
 backend_test.exclude(r'[a-z,_]*loop16_seq_none_[a-z,_]*')