Fix (tests): qonnx moved in brevitas_ort

Xilinx · Nov 14, 2023 · b3b9941 · b3b9941
1 parent baccb2c
commit b3b9941
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 132 deletions.
diff --git a/requirements/requirements-ort-integration.txt b/requirements/requirements-ort-integration.txt
@@ -1,3 +1,4 @@
 onnx
 onnxoptimizer
 onnxruntime
+qonnx
diff --git a/tests/brevitas_finn/brevitas/test_wbiol.py b/tests/brevitas_finn/brevitas/test_wbiol.py
diff --git a/tests/brevitas_ort/common.py b/tests/brevitas_ort/common.py
@@ -4,6 +4,9 @@
 import numpy as np
 import onnxruntime as ort
 import pytest
+from qonnx.core.modelwrapper import ModelWrapper
+import qonnx.core.onnx_exec as oxe
+from qonnx.transformation.infer_shapes import InferShapes
 import torch
 
 from brevitas.export import export_onnx_qcdq
@@ -17,10 +20,10 @@
 from brevitas.nn import QuantLinear
 from brevitas.nn import QuantLSTM
 from brevitas.nn import TruncAvgPool2d
-from brevitas.quant.fixed_point import Int8AccumulatorAwareWeightQuant
 from brevitas.quant.fixed_point import Int8ActPerTensorFixedPoint
 from brevitas.quant.fixed_point import Int8WeightPerChannelFixedPoint
 from brevitas.quant.fixed_point import Int8WeightPerTensorFixedPoint
+from brevitas.quant.scaled_int import Int8AccumulatorAwareWeightQuant
 from brevitas.quant.scaled_int import Int8ActPerTensorFloat
 from brevitas.quant.scaled_int import Int8WeightPerChannelFloat
 from brevitas.quant.scaled_int import Int8WeightPerTensorFloat
@@ -105,31 +108,39 @@ def is_brevitas_ort_close(
     input_t = torch.from_numpy(np_input)
     brevitas_output = model(input_t)
 
-    if export_type == 'qop':
-        export_onnx_qop(model, input_t, export_path=export_name)
-        brevitas_output = brevitas_output.int(float_datatype=False)
-    elif export_type == 'qcdq':
-        export_onnx_qcdq(model, input_t, export_path=export_name)
-    elif export_type == 'qcdq_opset14':
-        export_onnx_qcdq(model, input_t, opset_version=14, export_path=export_name)
-    elif export_type == 'qonnx_opset14':
-        export_qonnx(model, input_t, opset_version=14, export_path=export_name)
-    else:
-        raise RuntimeError(f"Export type {export_type} not recognized.")
-
     if tolerance is not None and export_type == 'qcdq':
         tolerance = tolerance * brevitas_output.scale  # Float Output, tolerance is +/- output scale
 
-    ort_output = compute_ort(export_name, np_input)
+    if export_type == 'qonnx':
+        exported_model = export_qonnx(model, input_t, export_path=export_name)
+        exported_model = ModelWrapper(exported_model)
+        exported_model = exported_model.transform(InferShapes())
+        idict = {exported_model.graph.input[0].name: np_input}
+        odict = oxe.execute_onnx(exported_model, idict, True)
+        ort_output = odict[exported_model.graph.output[0].name]
+    else:
+        if export_type == 'qop':
+            export_onnx_qop(model, input_t, export_path=export_name)
+            brevitas_output = brevitas_output.int(float_datatype=False)
+        elif export_type == 'qcdq':
+            export_onnx_qcdq(model, input_t, export_path=export_name)
+        elif export_type == 'qcdq_opset14':
+            export_onnx_qcdq(model, input_t, opset_version=14, export_path=export_name)
+        elif export_type == 'qonnx_opset14':
+            export_qonnx(model, input_t, opset_version=14, export_path=export_name)
+        else:
+            raise RuntimeError(f"Export type {export_type} not recognized.")
+
+        ort_output = compute_ort(export_name, np_input)
 
     if first_output_only:
-        if isinstance(ort_output, tuple):
+        if isinstance(ort_output, (tuple, list)):
             ort_output = ort_output[0]
         if isinstance(brevitas_output, tuple):
             brevitas_output = brevitas_output[0]
 
     # make sure we are not comparing 0s
-    if ort_output == 0 and (brevitas_output == 0).all():
+    if (ort_output == 0).all() and (brevitas_output == 0).all():
         pytest.skip("Skip testing against all 0s.")
 
     return recursive_allclose(ort_output, brevitas_output, tolerance)

diff --git a/tests/brevitas_ort/test_quant_module.py b/tests/brevitas_ort/test_quant_module.py
@@ -3,6 +3,7 @@
 
 from functools import reduce
 from operator import mul
+import os
 
 import pytest
 from pytest_cases import get_case_id
@@ -18,7 +19,7 @@
 
 
 @parametrize_with_cases('model', cases=QuantWBIOLCases)
-@pytest.mark.parametrize('export_type', ['qcdq', 'qop'])
+@pytest.mark.parametrize('export_type', ['qcdq', 'qonnx', 'qop'])
 @requires_pt_ge('1.8.1')
 def test_ort_wbiol(model, export_type, current_cases):
     cases_generator_func = current_cases['model'][1]