Merge pull request #332 from lucienwang1009/onnx_attr_check

Fetch valid attributes from ONNX dynamically

Author: nbcsm

.gitignore

@@ -1,4 +1,4 @@
-.coverage
+.coverage*
 *.pyc
 .idea
 build

tf2onnx/__init__.py

@@ -7,7 +7,7 @@
 from __future__ import unicode_literals
 
 
-__all__ = ["utils", "graph_matcher", "graph", "tfonnx", "shape_inference"]
+__all__ = ["utils", "graph_matcher", "graph", "tfonnx", "shape_inference", "schemas"]
 
 from .version import version as __version__
-from tf2onnx import tfonnx, utils, graph, graph_matcher, shape_inference  # pylint: disable=wrong-import-order
+from tf2onnx import tfonnx, utils, graph, graph_matcher, shape_inference, schemas  # pylint: disable=wrong-import-order

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.schemas import get_schema
 
 
 # todo(pengwa): remove protected-access later
@@ -72,8 +73,10 @@ def attr(self):
     def attr_onnx(self):
         onnx_attrs = {}
         for a in self._attr.values():
-            if a.name in utils.ONNX_VALID_ATTRIBUTES:
-                onnx_attrs[a.name] = a
+            schema = get_schema(self.type, self.graph.opset)
+            if schema:
+                if schema.has_attribute(a.name):
+                    onnx_attrs[a.name] = a
         return onnx_attrs
 
     @property

tf2onnx/schemas.py

@@ -0,0 +1,92 @@
+# 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
+
+
+ONNX_DOMAIN = ""
+
+
+class OnnxOpSchema(object):
+    """Wrapper for Onnx schema."""
+
+    def __init__(self, name, domain, since_version, attributes):
+        """Create a Onnx schema
+        Args:
+            name (str): op name
+            attributes (List[str]): valid attributes
+            domain (str): default value "" means it's Onnx domain
+            since_version (int): opset version, default is 1
+        """
+        self._name = name
+        self._domain = domain
+        self._attributes = attributes
+        self._since_version = since_version
+
+    @property
+    def attributes(self):
+        return self._attributes
+
+    @property
+    def domain(self):
+        return self._domain
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def since_version(self):
+        return self._since_version
+
+    @staticmethod
+    def from_onnx_schema(onnx_schema):
+        name = onnx_schema.name
+        domain = onnx_schema.domain
+        since_version = int(onnx_schema.since_version)
+        attributes = onnx_schema.attributes
+        return OnnxOpSchema(name, domain, since_version, attributes)
+
+    def has_attribute(self, attr):
+        return attr in self.attributes
+
+
+def _register_all_schemas_with_history():
+    """Register all schemas with history"""
+    onnx_schemas = defs.get_all_schemas_with_history()
+    name_domain_version_schema_map = defaultdict(lambda: defaultdict(dict))
+    for s in onnx_schemas:
+        schema = OnnxOpSchema.from_onnx_schema(s)
+        name_domain_version_schema_map[schema.name][schema.domain][schema.since_version] = schema
+
+    ordered_map = defaultdict(lambda: defaultdict(OrderedDict))
+    for name, domain_version_schema_map in name_domain_version_schema_map.items():
+        for domain, version_schema_map in domain_version_schema_map.items():
+            ordered_map[name][domain] = OrderedDict(
+                sorted(version_schema_map.items(), key=lambda x: -x[0])
+            )
+    return ordered_map
+
+
+# format is <OpName, <Domain, <SinceVersion, OpSchema>>>
+# SinceVersion is sorted from high to low
+_schemas = _register_all_schemas_with_history()
+
+
+def get_schema(name, max_inclusive_opset_version, domain=ONNX_DOMAIN):
+    """Get schema by name within specific version."""
+    domain_version_schema_map = _schemas[name]
+    version_schema_map = domain_version_schema_map[domain]
+    for version, schema in version_schema_map.items():
+        if version <= max_inclusive_opset_version:
+            return schema
+    return None

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