Skip to content
Snippets Groups Projects
Commit 255dec0b authored by Hendrik Borras's avatar Hendrik Borras
Browse files

Added QONNX_export test to test_brevitas_act_export_relu test.

parent a9a70b98
No related branches found
No related tags found
No related merge requests found
......@@ -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 QuantReLU
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_relu_act_export.onnx"
......@@ -50,7 +53,8 @@ export_onnx_path = "test_brevitas_relu_act_export.onnx"
@pytest.mark.parametrize(
"scaling_impl_type", [ScalingImplType.CONST, ScalingImplType.PARAMETER]
)
def test_brevitas_act_export_relu(abits, max_val, scaling_impl_type):
@pytest.mark.parametrize("QONNX_export", [False, True])
def test_brevitas_act_export_relu(abits, max_val, scaling_impl_type, QONNX_export):
min_val = -1.0
ishape = (1, 15)
......@@ -71,8 +75,15 @@ 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(
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment