Merge pull request #372 from nbcsm/data_format

node get/set data_format via attribute internally

Author: nbcsm

tests/test_internals.py

@@ -203,6 +203,21 @@ def test_shape_utils(self):
         self.assertFalse(utils.are_shapes_equal(None, []))
         self.assertTrue(utils.are_shapes_equal([1, 2, 3], (1, 2, 3)))
 
+    def test_data_format(self):
+        n1 = helper.make_node("Conv", ["X", "W"], ["Y"], name="n1", data_format="NHWC")
+        graph_proto = helper.make_graph(
+            nodes=[n1],
+            name="test",
+            inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 2]),
+                    helper.make_tensor_value_info("W", TensorProto.FLOAT, [2, 2])],
+            outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 2])],
+            initializer=[]
+        )
+        g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
+        n = g.get_node_by_name("n1")
+        self.assertEqual(n.data_format, "NHWC")
+        self.assertTrue(n.is_nhwc())
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/graph.py

@@ -46,10 +46,6 @@ def __init__(self, node, graph, skip_conversion=False):
         # dict to original attributes
         for a in node.attribute:
             self._attr[a.name] = a
-
-        self.data_format = self.get_attr("data_format")
-        if self.data_format:
-            self.data_format = self.data_format.s.decode("utf-8")
         self._skip_conversion = skip_conversion
 
     @property
@@ -123,6 +119,16 @@ def domain(self, val):
         """Set Op type."""
         self._op.domain = val
 
+    @property
+    def data_format(self):
+        """Return data_format."""
+        return self.get_attr_str("data_format")
+
+    @data_format.setter
+    def data_format(self, val):
+        """Set data_format."""
+        self.set_attr("data_format", val)
+
     def is_nhwc(self):
         """Return True if node is in NCHW format."""
         return self.data_format == "NHWC"
@@ -141,17 +147,22 @@ def __repr__(self):
         return "<onnx op type='%s' name=%s>" % (self.type, self._op.name)
 
     def get_attr(self, name, default=None):
-        """Get attribute map."""
+        """Get raw attribute value."""
         attr = self.attr.get(name, default)
         return attr
 
     def get_attr_int(self, name):
-        """Get attribute map."""
-        attr = self.attr.get(name)
+        """Get attribute value as int."""
+        attr = self.get_attr(name)
         utils.make_sure(attr is not None, "attribute %s is None", name)
         attr = attr.i
         return attr
 
+    def get_attr_str(self, name, encoding="utf-8"):
+        """Get attribute value as string."""
+        attr = self.get_attr(name)
+        return attr.s.decode(encoding) if attr else None
+
     def set_attr(self, name, value):
         self.attr[name] = helper.make_attribute(name, value)