fetch valid attributes from ONNX dynamically

Author: wayuanho

tf2onnx/graph.py

@@ -19,6 +19,7 @@
 from tf2onnx import utils, __version__
 from tf2onnx.utils import port_name, find_opset
 from tf2onnx.optimizer.transpose_optimizer import TransposeOptimizer
+from tf2onnx.schema import ONNXSchema
 
 
 # todo(pengwa): remove protected-access later
@@ -72,7 +73,7 @@ def attr(self):
     def attr_onnx(self):
         onnx_attrs = {}
         for a in self._attr.values():
-            if a.name in utils.ONNX_VALID_ATTRIBUTES:
+            if a.name in ONNXSchema.get_attribute(self.type, self.graph.opset):
                 onnx_attrs[a.name] = a
         return onnx_attrs
 

tf2onnx/schema.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.schema
+"""
+
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from collections import defaultdict, OrderedDict
+from onnx import defs
+
+def version_mapping(schemas):
+    """Obtain version to schema mapping."""
+    schemas_with_version = defaultdict(dict)
+    for s in schemas:
+        schemas_with_version[s.name][int(s.since_version)] = s
+    for name, s in schemas_with_version.items():
+        schemas_with_version[name] = OrderedDict(
+            sorted(schemas_with_version[name].items(), key=lambda x: x[0])
+        )
+    return schemas_with_version
+
+
+class ONNXSchema:
+    """Wrapper for ONNX schema"""
+
+    all_schemas = defs.get_all_schemas_with_history()
+    schemas_with_version = version_mapping(all_schemas)
+
+    @staticmethod
+    def get_schema(name, version):
+        """Get schema by name within specific version."""
+        if name not in ONNXSchema.schemas_with_version:
+            return None
+        if version < 1:
+            return None
+        schemas = ONNXSchema.schemas_with_version[name]
+        versions = list(schemas.keys())
+        for i, v in enumerate(versions):
+            if version < v:
+                return schemas[versions[i-1]]
+        return schemas[versions[-1]]
+
+    @staticmethod
+    def get_attribute(name, version):
+        """Get valid attributes by op's name and specific version"""
+        schema = ONNXSchema.get_schema(name, version)
+        if not schema:
+            return {}
+        return schema.attributes

tf2onnx/utils.py

@@ -76,21 +76,6 @@
 
 ONNX_UNKNOWN_DIMENSION = -1
 
-#
-# attributes onnx understands. Everything else coming from tensorflow
-# will be ignored.
-#
-ONNX_VALID_ATTRIBUTES = {
-    'p', 'bias', 'axes', 'pads', 'mean', 'activation_beta', 'spatial_scale', 'broadcast', 'pooled_shape', 'high',
-    'activation_alpha', 'is_test', 'hidden_size', 'activations', 'beta', 'input_as_shape', 'drop_states', 'alpha',
-    'momentum', 'scale', 'axis', 'dilations', 'transB', 'axis_w', 'blocksize', 'output_sequence', 'mode', 'perm',
-    'min', 'seed', 'ends', 'paddings', 'to', 'gamma', 'width_scale', 'normalize_variance', 'group', 'ratio', 'values',
-    'dtype', 'output_shape', 'spatial', 'split', 'input_forget', 'keepdims', 'transA', 'auto_pad', 'border', 'low',
-    'linear_before_reset', 'height_scale', 'output_padding', 'shape', 'kernel_shape', 'epsilon', 'size', 'starts',
-    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k', 'sample_size',
-    'blocksize', 'epsilon', 'momentum', 'body', 'directions', 'num_scan_inputs', 'then_branch', 'else_branch'
-}
-
 # index for internally generated names
 INTERNAL_NAME = 1