def test_brevitas_fc_onnx_export_and_exec(size, wbits, abits, pretrained): if size == "LFC" and wbits == 2 and abits == 2: pytest.skip(f"No LFC_{MAX_WBITS}W{MAX_ABITS}A present.") if wbits > abits: pytest.skip("No wbits > abits cases.") nname = f"{size}_{wbits}W{abits}A" finn_onnx = nname + ".onnx" fc, _ = model_with_cfg(nname.lower(), pretrained=pretrained) FINNManager.export_onnx(fc, FC_INPUT_SIZE, finn_onnx) model = ModelWrapper(finn_onnx) model = model.transform(GiveUniqueNodeNames()) model = model.transform(DoubleToSingleFloat()) model = model.transform(InferShapes()) model = model.transform(FoldConstants()) model = model.transform(RemoveStaticGraphInputs()) # load a random test vector input_tensor = np.random.uniform(MIN_INP_VAL, MAX_INP_VAL, size=FC_INPUT_SIZE).astype(np.float32) # run using FINN-based execution input_dict = {"0": input_tensor} output_dict = oxe.execute_onnx(model, input_dict) produced = output_dict[list(output_dict.keys())[0]] # run using PyTorch/Brevitas input_tensor = torch.from_numpy(input_tensor).float() # do forward pass in PyTorch/Brevitas expected = fc.forward(input_tensor).detach().numpy() assert np.isclose(produced, expected, atol=ATOL).all()
def test_brevitas_cnv_onnx_export_and_exec(wbits, abits, pretrained): if wbits > abits: pytest.skip("No wbits > abits cases.") nname = f"CNV_{wbits}W{abits}A" finn_onnx = nname + ".onnx" cnv, _ = model_with_cfg(nname.lower(), pretrained=pretrained) cnv.eval() # load a random int test vector input_a = np.random.randint(MIN_INP_VAL, MAX_INP_VAL, size=CNV_INPUT_SIZE).astype(np.float32) scale = 1. / 255 input_t = torch.from_numpy(input_a * scale) input_qt = QuantTensor( input_t, scale=torch.tensor(scale), bit_width=torch.tensor(8.0), signed=False) FINNManager.export(cnv, export_path=finn_onnx, input_t=input_qt) model = ModelWrapper(finn_onnx) model = model.transform(GiveUniqueNodeNames()) model = model.transform(DoubleToSingleFloat()) model = model.transform(InferShapes()) model = model.transform(FoldConstants()) model = model.transform(RemoveStaticGraphInputs()) # run using FINN-based execution input_dict = {"0": input_a} output_dict = oxe.execute_onnx(model, input_dict) produced = output_dict[list(output_dict.keys())[0]] # do forward pass in PyTorch/Brevitas expected = cnv(input_t).detach().numpy() assert np.isclose(produced, expected, atol=ATOL).all()
def test_debug_finn_onnx_export(): model, cfg = model_with_cfg(REF_MODEL, pretrained=False) debug_hook = enable_debug(model) input_tensor = torch.randn(1, 3, 32, 32) export_finn_onnx(model, input_shape=input_tensor.shape, export_path='debug.onnx') model(input_tensor) assert debug_hook.values
def test_brevitas_cnv_jit_trace(wbits, abits): if wbits > abits: pytest.skip("No wbits > abits cases.") nname = f"CNV_{wbits}W{abits}A" cnv, _ = model_with_cfg(nname.lower(), pretrained=False) cnv.train(False) input_tensor = torch.randn(CNV_INPUT_SIZE) traced_model = jit_trace_patched(cnv, input_tensor) out_traced = traced_model(input_tensor) out = cnv(input_tensor) assert out.isclose(out_traced).all().item()
def test_brevitas_fc_jit_trace(size, wbits, abits): if size == "LFC" and wbits == 2 and abits == 2: pytest.skip(f"No LFC_{MAX_WBITS}W{MAX_ABITS}A present.") if wbits > abits: pytest.skip("No wbits > abits cases.") nname = f"{size}_{wbits}W{abits}A" fc, _ = model_with_cfg(nname.lower(), pretrained=False) fc.train(False) input_tensor = torch.randn(FC_INPUT_SIZE) traced_model = jit_trace_patched(fc, input_tensor) out_traced = traced_model(input_tensor) out = fc(input_tensor) assert out.isclose(out_traced).all().item()
def test_brevitas_cnv_jit_trace(wbits, abits): if wbits > abits: pytest.skip("No wbits > abits cases.") nname = f"CNV_{wbits}W{abits}A" cnv, _ = model_with_cfg(nname.lower(), pretrained=False) cnv.train(False) input_tensor = torch.randn(CNV_INPUT_SIZE) with ExitStack() as stack: for mgr in jit_patches_generator(): stack.enter_context(mgr) traced_model = torch.jit.trace(cnv, input_tensor) out_traced = traced_model(input_tensor) out = cnv(input_tensor) assert out.isclose(out_traced).all().item()