diff --git a/tests/brevitas/test_brevitas_scaled_QHardTanh_export.py b/tests/brevitas/test_brevitas_scaled_QHardTanh_export.py index 345fae872119c75aa8e85cb5812c94dfc15bad7f..c6da2e2e971ee97cb73243284920cc87e8b4d7bb 100644 --- a/tests/brevitas/test_brevitas_scaled_QHardTanh_export.py +++ b/tests/brevitas/test_brevitas_scaled_QHardTanh_export.py @@ -36,11 +36,14 @@ import torch from brevitas.core.quant import QuantType from brevitas.core.restrict_val import RestrictValueType from brevitas.core.scaling import ScalingImplType +from brevitas.export.onnx.generic.manager import BrevitasONNXManager from brevitas.nn import QuantHardTanh +from qonnx.util.cleanup import cleanup as qonnx_cleanup import finn.core.onnx_exec as oxe from finn.core.modelwrapper import ModelWrapper from finn.transformation.infer_shapes import InferShapes +from finn.transformation.qonnx.convert_qonnx_to_finn import ConvertQONNXtoFINN export_onnx_path = "test_brevitas_scaled_QHardTanh_export.onnx" @@ -52,8 +55,9 @@ export_onnx_path = "test_brevitas_scaled_QHardTanh_export.onnx" @pytest.mark.parametrize( "scaling_impl_type", [ScalingImplType.CONST, ScalingImplType.PARAMETER] ) +@pytest.mark.parametrize("QONNX_export", [False, True]) def test_brevitas_act_export_qhardtanh_scaled( - abits, narrow_range, min_val, max_val, scaling_impl_type + abits, narrow_range, min_val, max_val, scaling_impl_type, QONNX_export ): def get_quant_type(bit_width): if bit_width is None: @@ -84,8 +88,15 @@ tensor_quant.scaling_impl.learned_value": torch.tensor( ) } b_act.load_state_dict(checkpoint) - - bo.export_finn_onnx(b_act, ishape, export_onnx_path) + if QONNX_export: + m_path = export_onnx_path + BrevitasONNXManager.export(b_act, ishape, m_path) + qonnx_cleanup(m_path, out_file=m_path) + model = ModelWrapper(m_path) + model = model.transform(ConvertQONNXtoFINN()) + model.save(m_path) + else: + bo.export_finn_onnx(b_act, ishape, export_onnx_path) model = ModelWrapper(export_onnx_path) model = model.transform(InferShapes()) inp_tensor = np.random.uniform(low=min_val, high=max_val, size=ishape).astype(