def test_generic_quant_conv_export():
IN_SIZE = (2, IN_CH, IN_CH, IN_CH)
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = QuantConv2d(out_channels=OUT_CH,
in_channels=IN_CH,
bias=True,
kernel_size=3,
input_quant=Int8ActPerTensorFloat,
output_quant=Int8ActPerTensorFloat,
bias_quant=Int16Bias,
return_quant_tensor=False)
self.conv.weight.data.uniform_(-0.01, 0.01)
def forward(self, x):
return self.conv(x)
inp = torch.randn(IN_SIZE)
model = Model()
model(inp) # collect scale factors
model.eval()
BrevitasONNXManager.export(model,
input_t=inp,
export_path='./generic_quant_conv.onnx')
def test_generic_quant_tensor_export():
IN_SIZE = (2, IN_CH)
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.quant_inp = QuantIdentity(return_quant_tensor=True)
self.linear = QuantLinear(out_features=OUT_CH,
in_features=IN_CH,
bias=True,
output_quant=Int8ActPerTensorFloat,
bias_quant=Int16Bias,
return_quant_tensor=False)
self.linear.weight.data.uniform_(-0.01, 0.01)
def forward(self, x):
return self.linear(self.quant_inp(x))
inp = torch.randn(IN_SIZE)
model = Model()
model(inp) # collect scale factors
model.eval()
BrevitasONNXManager.export(model,
input_t=inp,
export_path='./generic_quant_tensor.onnx')
def test_export(self, topology, wbits, abits, QONNX_export):
if wbits > abits:
pytest.skip("No wbits > abits end2end network configs for now")
if topology == "lfc" and not (wbits == 1 and abits == 1):
pytest.skip("Skipping certain lfc configs")
(model,
ishape) = get_trained_network_and_ishape(topology, wbits, abits)
chkpt_name = get_checkpoint_name(topology, wbits, abits, QONNX_export,
"export")
if QONNX_export:
BrevitasONNXManager.export(model, ishape, chkpt_name)
qonnx_cleanup(chkpt_name, out_file=chkpt_name)
model = ModelWrapper(chkpt_name)
model = model.transform(ConvertQONNXtoFINN())
model.save(chkpt_name)
else:
bo.export_finn_onnx(model, ishape, chkpt_name)
nname = "%s_w%da%d" % (topology, wbits, abits)
update_dashboard_data(topology, wbits, abits, "network", nname)
dtstr = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
update_dashboard_data(topology, wbits, abits, "datetime", dtstr)
finn_commit = subprocess.check_output(["git", "rev-parse", "HEAD"],
cwd="/workspace/finn")
finn_commit = finn_commit.decode("utf-8").strip()
update_dashboard_data(topology, wbits, abits, "finn-commit",
finn_commit)
assert os.path.isfile(chkpt_name)
def test_brevitas_cnv_export_exec(wbits, abits, QONNX_export):
if wbits > abits:
pytest.skip("No wbits > abits cases at the moment")
cnv = get_test_model_trained("CNV", wbits, abits)
ishape = (1, 3, 32, 32)
if QONNX_export:
BrevitasONNXManager.export(cnv, ishape, export_onnx_path)
qonnx_cleanup(export_onnx_path, out_file=export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(ConvertQONNXtoFINN())
model.save(export_onnx_path)
else:
bo.export_finn_onnx(cnv, ishape, export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(GiveUniqueNodeNames())
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(RemoveStaticGraphInputs())
assert len(model.graph.input) == 1
assert len(model.graph.output) == 1
fn = pk.resource_filename("finn.qnn-data", "cifar10/cifar10-test-data-class3.npz")
input_tensor = np.load(fn)["arr_0"].astype(np.float32)
input_tensor = input_tensor / 255
assert input_tensor.shape == (1, 3, 32, 32)
# run using FINN-based execution
input_dict = {model.graph.input[0].name: input_tensor}
output_dict = oxe.execute_onnx(model, input_dict, True)
produced = output_dict[model.graph.output[0].name]
# do forward pass in PyTorch/Brevitas
input_tensor = torch.from_numpy(input_tensor).float()
expected = cnv.forward(input_tensor).detach().numpy()
assert np.isclose(produced, expected, atol=1e-3).all()
assert np.argmax(produced) == 3
os.remove(export_onnx_path)
def test_brevitas_QConv2d(dw, bias, in_channels, QONNX_export):
ishape = (1, 32, 111, 111)
if dw is True:
groups = in_channels
out_channels = in_channels
kernel_size = 3
padding = 1
stride = 1
w_shape = (32, 1, 3, 3)
else:
groups = 1
out_channels = 64
kernel_size = 1
padding = 0
stride = 1
w_shape = (64, 32, 1, 1)
b_conv = QuantConv2d(
in_channels=in_channels,
out_channels=out_channels,
groups=groups,
kernel_size=kernel_size,
padding=padding,
stride=stride,
bias=bias,
bias_quant_type=QuantType.FP,
weight_bit_width=4,
weight_quant_type=QuantType.INT,
weight_scaling_impl_type=ScalingImplType.STATS,
weight_scaling_stats_op=StatsOp.MAX,
weight_scaling_per_output_channel=True,
weight_restrict_scaling_type=RestrictValueType.LOG_FP,
weight_narrow_range=True,
weight_scaling_min_val=2e-16,
)
weight_tensor = gen_finn_dt_tensor(DataType["INT4"], w_shape)
b_conv.weight = torch.nn.Parameter(torch.from_numpy(weight_tensor).float())
b_conv.eval()
if QONNX_export:
m_path = export_onnx_path
BrevitasONNXManager.export(b_conv, 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_conv, ishape, export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(InferShapes())
inp_tensor = np.random.uniform(low=-1.0, high=1.0, size=ishape).astype(np.float32)
idict = {model.graph.input[0].name: inp_tensor}
odict = oxe.execute_onnx(model, idict, True)
produced = odict[model.graph.output[0].name]
inp_tensor = torch.from_numpy(inp_tensor).float()
expected = b_conv.forward(inp_tensor).detach().numpy()
assert np.isclose(produced, expected, atol=1e-3).all()
os.remove(export_onnx_path)
def test_generic_quant_avgpool_export_quant_input():
IN_SIZE = (2, OUT_CH, IN_CH, IN_CH)
inp = torch.randn(IN_SIZE)
inp_quant = QuantIdentity(return_quant_tensor=True)
model = QuantAvgPool2d(kernel_size=2, return_quant_tensor=False)
inp_quant(inp) # collect scale factors
inp_quant.eval()
model.eval()
BrevitasONNXManager.export(
model, input_t=inp_quant(inp), export_path='generic_quant_avgpool_quant_input.onnx')
def test_brevitas_act_export_relu(abits, max_val, scaling_impl_type,
QONNX_export):
min_val = -1.0
ishape = (1, 15)
b_act = QuantReLU(
bit_width=abits,
max_val=max_val,
scaling_impl_type=scaling_impl_type,
restrict_scaling_type=RestrictValueType.LOG_FP,
quant_type=QuantType.INT,
)
if scaling_impl_type == ScalingImplType.PARAMETER:
checkpoint = {
"act_quant_proxy.fused_activation_quant_proxy.tensor_quant.\
scaling_impl.learned_value":
torch.tensor(0.49).type(torch.FloatTensor)
}
b_act.load_state_dict(checkpoint)
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(np.float32)
idict = {model.graph.input[0].name: inp_tensor}
odict = oxe.execute_onnx(model, idict, True)
produced = odict[model.graph.output[0].name]
inp_tensor = torch.from_numpy(inp_tensor).float()
b_act.eval()
expected = b_act.forward(inp_tensor).detach().numpy()
if not np.isclose(produced, expected, atol=1e-3).all():
print(abits, max_val, scaling_impl_type)
print("scale: ",
b_act.quant_act_scale().type(torch.FloatTensor).detach())
if abits < 5:
print(
"thres:",
", ".join(["{:8.4f}".format(x)
for x in b_act.export_thres[0]]),
)
print("input:",
", ".join(["{:8.4f}".format(x) for x in inp_tensor[0]]))
print("prod :", ", ".join(["{:8.4f}".format(x) for x in produced[0]]))
print("expec:", ", ".join(["{:8.4f}".format(x) for x in expected[0]]))
assert np.isclose(produced, expected, atol=1e-3).all()
os.remove(export_onnx_path)
def test_brevitas_act_export_qhardtanh_nonscaled(abits, narrow_range, max_val,
QONNX_export):
def get_quant_type(bit_width):
if bit_width is None:
return QuantType.FP
elif bit_width == 1:
return QuantType.BINARY
else:
return QuantType.INT
act_quant_type = get_quant_type(abits)
min_val = -1.0
ishape = (1, 10)
b_act = QuantHardTanh(
bit_width=abits,
quant_type=act_quant_type,
max_val=max_val,
min_val=min_val,
restrict_scaling_type=RestrictValueType.LOG_FP,
scaling_impl_type=ScalingImplType.CONST,
narrow_range=narrow_range,
)
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(np.float32)
idict = {model.graph.input[0].name: inp_tensor}
odict = oxe.execute_onnx(model, idict, True)
produced = odict[model.graph.output[0].name]
inp_tensor = torch.from_numpy(inp_tensor).float()
expected = b_act.forward(inp_tensor).detach().numpy()
assert np.isclose(produced, expected, atol=1e-3).all()
os.remove(export_onnx_path)
def test_generic_quant_avgpool_export():
IN_SIZE = (2, OUT_CH, IN_CH, IN_CH)
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.inp_quant = QuantIdentity(return_quant_tensor=True)
self.pool = QuantAvgPool2d(kernel_size=2)
def forward(self, x):
return self.pool(self.inp_quant(x))
inp = torch.randn(IN_SIZE)
model = Model()
model(inp) # collect scale factors
model.eval()
BrevitasONNXManager.export(model,
input_t=inp,
export_path='./generic_quant_avgpool.onnx')
def test_brevitas_qlinear(
bias, out_features, in_features, w_bits, i_dtype, QONNX_export
):
i_shape = (1, in_features)
w_shape = (out_features, in_features)
b_linear = QuantLinear(
out_features=out_features,
in_features=in_features,
bias=bias,
bias_quant_type=QuantType.FP,
weight_bit_width=w_bits,
weight_quant_type=QuantType.INT,
weight_scaling_per_output_channel=True,
)
weight_tensor_fp = np.random.uniform(low=-1.0, high=1.0, size=w_shape).astype(
np.float32
)
b_linear.weight.data = torch.from_numpy(weight_tensor_fp)
b_linear.eval()
if QONNX_export:
m_path = export_onnx_path
BrevitasONNXManager.export(b_linear, i_shape, 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_linear, i_shape, export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(InferShapes())
inp_tensor = gen_finn_dt_tensor(i_dtype, i_shape)
idict = {model.graph.input[0].name: inp_tensor}
odict = oxe.execute_onnx(model, idict, True)
produced = odict[model.graph.output[0].name]
inp_tensor = torch.from_numpy(inp_tensor).float()
expected = b_linear.forward(inp_tensor).detach().numpy()
assert np.isclose(produced, expected, atol=1e-3).all()
os.remove(export_onnx_path)
def test_brevitas_fc_onnx_export_and_exec(size, wbits, abits, QONNX_export):
if size == "LFC" and wbits == 2 and abits == 2:
pytest.skip("No LFC-w2a2 present at the moment")
if wbits > abits:
pytest.skip("No wbits > abits cases at the moment")
nname = "%s_%dW%dA_QONNX-%d" % (size, wbits, abits, QONNX_export)
finn_onnx = export_onnx_path + "/%s.onnx" % nname
fc = get_test_model_trained(size, wbits, abits)
ishape = (1, 1, 28, 28)
if QONNX_export:
BrevitasONNXManager.export(fc, ishape, finn_onnx)
qonnx_cleanup(finn_onnx, out_file=finn_onnx)
model = ModelWrapper(finn_onnx)
model = model.transform(ConvertQONNXtoFINN())
model.save(finn_onnx)
else:
bo.export_finn_onnx(fc, ishape, finn_onnx)
model = ModelWrapper(finn_onnx)
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(RemoveStaticGraphInputs())
assert len(model.graph.input) == 1
assert len(model.graph.output) == 1
# load one of the test vectors
raw_i = get_data("finn.data", "onnx/mnist-conv/test_data_set_0/input_0.pb")
input_tensor = onnx.load_tensor_from_string(raw_i)
# run using FINN-based execution
input_dict = {model.graph.input[0].name: nph.to_array(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(nph.to_array(input_tensor)).float()
assert input_tensor.shape == (1, 1, 28, 28)
# do forward pass in PyTorch/Brevitas
expected = fc.forward(input_tensor).detach().numpy()
assert np.isclose(produced, expected, atol=1e-3).all()
def test_brevitas_avg_pool_export(
kernel_size,
stride,
signed,
bit_width,
input_bit_width,
channels,
idim,
QONNX_export,
):
export_onnx_path = base_export_onnx_path.replace(
".onnx", f"test_QONNX-{QONNX_export}.onnx"
)
quant_avgpool = QuantAvgPool2d(
kernel_size=kernel_size,
stride=stride,
bit_width=bit_width,
return_quant_tensor=False,
)
quant_avgpool.eval()
# determine input
prefix = "INT" if signed else "UINT"
dt_name = prefix + str(input_bit_width)
dtype = DataType[dt_name]
input_shape = (1, channels, idim, idim)
input_array = gen_finn_dt_tensor(dtype, input_shape)
# Brevitas QuantAvgPool layers need QuantTensors to export correctly
# which requires setting up a QuantTensor instance with the scale
# factor, zero point, bitwidth and signedness
scale_array = np.ones((1, channels, 1, 1)).astype(np.float32)
scale_array *= 0.5
input_tensor = torch.from_numpy(input_array * scale_array).float()
scale_tensor = torch.from_numpy(scale_array).float()
zp = torch.tensor(0.0)
input_quant_tensor = QuantTensor(
input_tensor, scale_tensor, zp, input_bit_width, signed, training=False
)
# export
if QONNX_export:
BrevitasONNXManager.export(
quant_avgpool,
export_path=export_onnx_path,
input_t=input_quant_tensor,
)
model = ModelWrapper(export_onnx_path)
# Statically set the additional inputs generated by the BrevitasONNXManager
model.graph.input.remove(model.graph.input[3])
model.graph.input.remove(model.graph.input[2])
model.graph.input.remove(model.graph.input[1])
model.set_initializer("1", scale_array)
model.set_initializer("2", np.array(0.0).astype(np.float32))
model.set_initializer("3", np.array(input_bit_width).astype(np.float32))
model.save(export_onnx_path)
qonnx_cleanup(export_onnx_path, out_file=export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(ConvertQONNXtoFINN())
model.save(export_onnx_path)
else:
FINNManager.export(
quant_avgpool, export_path=export_onnx_path, input_t=input_quant_tensor
)
model = ModelWrapper(export_onnx_path)
model = model.transform(InferShapes())
model = model.transform(InferDataTypes())
# reference brevitas output
ref_output_array = quant_avgpool(input_quant_tensor).detach().numpy()
# finn output
if QONNX_export:
# Manually apply the Quant tensor scaling for QONNX
idict = {model.graph.input[0].name: input_array * scale_array}
else:
idict = {model.graph.input[0].name: input_array}
odict = oxe.execute_onnx(model, idict, True)
finn_output = odict[model.graph.output[0].name]
# compare outputs
assert np.isclose(ref_output_array, finn_output).all()
# cleanup
# assert False
os.remove(export_onnx_path)
def test_end2end_cybsec_mlp_export(QONNX_export):
assets_dir = pk.resource_filename("finn.qnn-data", "cybsec-mlp/")
# load up trained net in Brevitas
input_size = 593
hidden1 = 64
hidden2 = 64
hidden3 = 64
weight_bit_width = 2
act_bit_width = 2
num_classes = 1
model = nn.Sequential(
QuantLinear(input_size,
hidden1,
bias=True,
weight_bit_width=weight_bit_width),
nn.BatchNorm1d(hidden1),
nn.Dropout(0.5),
QuantReLU(bit_width=act_bit_width),
QuantLinear(hidden1,
hidden2,
bias=True,
weight_bit_width=weight_bit_width),
nn.BatchNorm1d(hidden2),
nn.Dropout(0.5),
QuantReLU(bit_width=act_bit_width),
QuantLinear(hidden2,
hidden3,
bias=True,
weight_bit_width=weight_bit_width),
nn.BatchNorm1d(hidden3),
nn.Dropout(0.5),
QuantReLU(bit_width=act_bit_width),
QuantLinear(hidden3,
num_classes,
bias=True,
weight_bit_width=weight_bit_width),
)
trained_state_dict = torch.load(assets_dir +
"/state_dict.pth")["models_state_dict"][0]
model.load_state_dict(trained_state_dict, strict=False)
W_orig = model[0].weight.data.detach().numpy()
# pad the second (593-sized) dimensions with 7 zeroes at the end
W_new = np.pad(W_orig, [(0, 0), (0, 7)])
model[0].weight.data = torch.from_numpy(W_new)
model_for_export = CybSecMLPForExport(model)
export_onnx_path = get_checkpoint_name("export", QONNX_export)
input_shape = (1, 600)
# create a QuantTensor instance to mark the input as bipolar during export
input_a = np.random.randint(0, 1, size=input_shape).astype(np.float32)
input_a = 2 * input_a - 1
scale = 1.0
input_t = torch.from_numpy(input_a * scale)
input_qt = QuantTensor(input_t,
scale=torch.tensor(scale),
bit_width=torch.tensor(1.0),
signed=True)
if QONNX_export:
# With the BrevitasONNXManager we need to manually set
# the FINN DataType at the input
BrevitasONNXManager.export(model_for_export,
input_shape,
export_path=export_onnx_path)
model = ModelWrapper(export_onnx_path)
model.set_tensor_datatype(model.graph.input[0].name,
DataType["BIPOLAR"])
model.save(export_onnx_path)
qonnx_cleanup(export_onnx_path, out_file=export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(ConvertQONNXtoFINN())
model.save(export_onnx_path)
else:
bo.export_finn_onnx(model_for_export,
export_path=export_onnx_path,
input_t=input_qt)
assert os.path.isfile(export_onnx_path)
# fix input datatype
finn_model = ModelWrapper(export_onnx_path)
finnonnx_in_tensor_name = finn_model.graph.input[0].name
assert tuple(finn_model.get_tensor_shape(finnonnx_in_tensor_name)) == (1,
600)
# verify a few exported ops
if QONNX_export:
# The first "Mul" node doesn't exist in the QONNX export,
# because the QuantTensor scale is not exported.
# However, this node would have been unity scale anyways and
# the models are still equivalent.
assert finn_model.graph.node[0].op_type == "Add"
assert finn_model.graph.node[1].op_type == "Div"
assert finn_model.graph.node[2].op_type == "MatMul"
assert finn_model.graph.node[-1].op_type == "MultiThreshold"
else:
assert finn_model.graph.node[0].op_type == "Mul"
assert finn_model.get_initializer(
finn_model.graph.node[0].input[1]) == 1.0
assert finn_model.graph.node[1].op_type == "Add"
assert finn_model.graph.node[2].op_type == "Div"
assert finn_model.graph.node[3].op_type == "MatMul"
assert finn_model.graph.node[-1].op_type == "MultiThreshold"
# verify datatypes on some tensors
assert (finn_model.get_tensor_datatype(finnonnx_in_tensor_name) ==
DataType["BIPOLAR"])
first_matmul_w_name = finn_model.get_nodes_by_op_type("MatMul")[0].input[1]
assert finn_model.get_tensor_datatype(
first_matmul_w_name) == DataType["INT2"]
def test_brevitas_debug(QONNX_export, QONNX_FINN_conversion):
if (not QONNX_export) and QONNX_FINN_conversion:
pytest.skip(
"This test configuration is not valid and is thus skipped.")
finn_onnx = "test_brevitas_debug.onnx"
fc = get_test_model_trained("TFC", 2, 2)
ishape = (1, 1, 28, 28)
if QONNX_export:
dbg_hook = bo.enable_debug(fc, proxy_level=True)
BrevitasONNXManager.export(fc, ishape, finn_onnx)
# DebugMarkers have the brevitas.onnx domain, so that needs adjusting
model = ModelWrapper(finn_onnx)
dbg_nodes = model.get_nodes_by_op_type("DebugMarker")
for dbg_node in dbg_nodes:
dbg_node.domain = "finn.custom_op.general"
model.save(finn_onnx)
qonnx_cleanup(finn_onnx, out_file=finn_onnx)
if QONNX_FINN_conversion:
model = ModelWrapper(finn_onnx)
model = model.transform(ConvertQONNXtoFINN())
model.save(finn_onnx)
else:
dbg_hook = bo.enable_debug(fc)
bo.export_finn_onnx(fc, ishape, finn_onnx)
model = ModelWrapper(finn_onnx)
# DebugMarkers have the brevitas.onnx domain, so that needs adjusting
# ToDo: We should probably have transformation pass, which does this
# domain conversion for us?
dbg_nodes = model.get_nodes_by_op_type("DebugMarker")
for dbg_node in dbg_nodes:
dbg_node.domain = "finn.custom_op.general"
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(RemoveStaticGraphInputs())
model.save(finn_onnx)
model = ModelWrapper(finn_onnx)
assert len(model.graph.input) == 1
assert len(model.graph.output) == 1
# load one of the test vectors
raw_i = get_data("finn.data", "onnx/mnist-conv/test_data_set_0/input_0.pb")
input_tensor = onnx.load_tensor_from_string(raw_i)
# run using FINN-based execution
input_dict = {model.graph.input[0].name: nph.to_array(input_tensor)}
output_dict = oxe.execute_onnx(model,
input_dict,
return_full_exec_context=True)
produced = output_dict[model.graph.output[0].name]
# run using PyTorch/Brevitas
input_tensor = torch.from_numpy(nph.to_array(input_tensor)).float()
assert input_tensor.shape == (1, 1, 28, 28)
# do forward pass in PyTorch/Brevitas
expected = fc.forward(input_tensor).detach().numpy()
assert np.isclose(produced, expected, atol=1e-3).all()
# check all tensors at debug markers
names_brevitas = set(dbg_hook.values.keys())
names_finn = set(output_dict.keys())
names_common = names_brevitas.intersection(names_finn)
# The different exports return debug markers in different numbers and places
print(len(names_common))
if QONNX_export and not QONNX_FINN_conversion:
assert len(names_common) == 12
elif QONNX_export and QONNX_FINN_conversion:
assert len(names_common) == 8
else:
assert len(names_common) == 16
for dbg_name in names_common:
if QONNX_export:
tensor_pytorch = dbg_hook.values[dbg_name].value.detach().numpy()
else:
tensor_pytorch = dbg_hook.values[dbg_name].detach().numpy()
tensor_finn = output_dict[dbg_name]
assert np.isclose(tensor_finn, tensor_pytorch, atol=1e-5).all()
os.remove(finn_onnx)