Support infer_datatype for flatten layer (#30)

* Support infer_datatype for flatten layer * [InferDT] add more identity op types for datatype inference * [Lint] fix linting issues Co-authored-by: Yaman Umuroglu <[email protected]>
Xilinx · May 17, 2021 · ac0b86a · ac0b86a
1 parent 2c08044
commit ac0b86a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/src/finn/transformation/infer_datatypes.py b/src/finn/transformation/infer_datatypes.py
@@ -35,7 +35,14 @@
 def _infer_node_datatype(model, node):
     """Infer output datatype(s) for a particular node. Returns True if any
     changes were made."""
-    dt_identity_optypes = ["Reshape", "Transpose"]
+    dt_identity_optypes = [
+        "Reshape",
+        "Transpose",
+        "Flatten",
+        "Slice",
+        "Gather",
+        "Identity",
+    ]
     idtypes = list(map(lambda x: model.get_tensor_datatype(x), node.input))
     odtypes = list(map(lambda x: model.get_tensor_datatype(x), node.output))
     op_type = node.op_type

diff --git a/tests/transformation/test_4d_conversion.py b/tests/transformation/test_4d_conversion.py
@@ -24,7 +24,7 @@ def generate_random_input(model):
 
 
 def set_all_initializers(model):
-    """ Sets all initializers of the graph to a random value. """
+    """Sets all initializers of the graph to a random value."""
     for n in model.graph.node:
         if len(n.input) > 1:
             init_name = n.input[1]

diff --git a/tests/transformation/test_batchnorm_to_affine.py b/tests/transformation/test_batchnorm_to_affine.py
@@ -70,7 +70,7 @@ def test_batchnorm_to_affine_shufflenet():
 
 @pytest.mark.parametrize("epsilon", [0.0, 0.00001, 0.001])
 def test_batchnorm_to_affine_epsilon(epsilon):
-    """ Dummy batchnorm node to test out the epsilon attribute. """
+    """Dummy batchnorm node to test out the epsilon attribute."""
 
     batchnorm_node = onnx.helper.make_node(
         "BatchNormalization",