Add option to disable scaled-int datatypes in InferDataTypes

src/qonnx/transformation/infer_datatypes.py

@@ -52,7 +52,7 @@ def infer_mac_result_dtype(idtypes, odtype_orig, possible_negation):
     return ret
 
 
-def _infer_node_datatype(model, node):
+def _infer_node_datatype(model, node, allow_scaledint_dtypes):
     """Infer output datatype(s) for a particular node. Returns True if any
     changes were made."""
     dt_identity_optypes = [
@@ -137,6 +137,11 @@ def _infer_node_datatype(model, node):
                     model.set_tensor_datatype(o, odtype)
                 else:
                     model.set_tensor_datatype(o, DataType["FLOAT32"])
+    # if scaled-int dtype inference is disabled, replace those with FLOAT32
+    if not allow_scaledint_dtypes:
+        for out in node.output:
+            if "SCALEDINT" in model.get_tensor_datatype(out).get_canonical_name():
+                model.set_tensor_datatype(out, DataType["FLOAT32"])
     # compare old and new output dtypes to see if anything changed
     new_odtypes = list(map(lambda x: model.get_tensor_datatype(x), node.output))
     graph_modified = new_odtypes != odtypes
@@ -147,9 +152,13 @@ class InferDataTypes(Transformation):
     """Infer QONNX DataType info for all intermediate/output tensors based on
     inputs and node type."""
 
+    def __init__(self, allow_scaledint_dtypes=False):
+        super().__init__()
+        self.allow_scaledint_dtypes = allow_scaledint_dtypes
+
     def apply(self, model):
         graph = model.graph
         graph_modified = False
         for node in graph.node:
-            graph_modified |= _infer_node_datatype(model, node)
+            graph_modified |= _infer_node_datatype(model, node, self.allow_scaledint_dtypes)
         return (model, graph_modified)

src/qonnx/util/inference_cost.py

@@ -114,11 +114,11 @@ def inference_cost(
             qnt_node.domain = "qonnx.custom_op.general"
         model = model.transform(InferShapes())
         model = model.transform(GiveUniqueParameterTensors())
-        model = model.transform(InferDataTypes())
+        model = model.transform(InferDataTypes(allow_scaledint_dtypes=True))
         model = model.transform(FoldConstants(exclude_op_types=[]))
         model = model.transform(RemoveUnusedTensors())
         model = model.transform(RemoveStaticGraphInputs())
-        model = model.transform(InferDataTypes())
+        model = model.transform(InferDataTypes(allow_scaledint_dtypes=True))
     model = model.transform(GiveUniqueNodeNames())
     model = model.transform(GiveReadableTensorNames())
     if output_onnx is not None:

tests/transformation/test_infer_datatypes.py

@@ -34,6 +34,7 @@
 from qonnx.transformation.general import GiveReadableTensorNames, GiveUniqueNodeNames
 from qonnx.transformation.infer_datatypes import InferDataTypes, infer_mac_result_dtype
 from qonnx.transformation.infer_shapes import InferShapes
+from qonnx.util.test import download_model
 
 
 def test_infer_mac_dtype_result():
@@ -79,3 +80,17 @@ def test_infer_datatypes():
     assert model.get_tensor_datatype("Conv_0_out0") == DataType["INT32"]
     assert model.get_tensor_datatype("Relu_0_out0") == DataType["FLOAT32"]
     assert model.get_tensor_datatype("global_out") == DataType["FLOAT32"]
+
+
+def test_infer_datatypes_scaledint():
+    orig_model = download_model("FINN-CNV_W2A2", do_cleanup=True, return_modelwrapper=True)
+    model = orig_model.transform(InferDataTypes(allow_scaledint_dtypes=True))
+    assert model.get_tensor_datatype("Quant_9_out0") == DataType["SCALEDINT<8>"]
+    assert model.get_tensor_datatype("Conv_0_out0") == DataType["SCALEDINT<32>"]
+    model = orig_model.transform(InferDataTypes(allow_scaledint_dtypes=False))
+    assert model.get_tensor_datatype("Quant_9_out0") == DataType["FLOAT32"]
+    assert model.get_tensor_datatype("Conv_0_out0") == DataType["FLOAT32"]
+    # no dtypes should be inferred as SCALEDINT
+    for tname in model.get_all_tensor_names():
+        tensor_dt = model.get_tensor_datatype(tname)
+        assert not ("SCALEDINT" in tensor_dt.get_canonical_name())

tests/transformation/test_pruning.py

@@ -86,7 +86,7 @@ def test_pruning_mnv1():
     model = download_model("MobileNetv1-w4a4", return_modelwrapper=True)
     # mark input as scaled 8-bit to get correct inference cost
     model.set_tensor_datatype(model.graph.input[0].name, DataType["SCALEDINT<8>"])
-    model = model.transform(InferDataTypes())
+    model = model.transform(InferDataTypes(allow_scaledint_dtypes=True))
     # do cleanup including folding quantized weights
     model = cleanup_model(model, False)
     inp, golden = get_golden_in_and_output("MobileNetv1-w4a4")