def test_xnorpopcountmatmul():
M = 1
K = 3
N = 3
x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [M, K])
W = helper.make_tensor_value_info("W", TensorProto.FLOAT, [K, N])
out = helper.make_tensor_value_info("out", TensorProto.FLOAT, ["x", "y"])
node_def = helper.make_node("XnorPopcountMatMul", ["x", "W"], ["out"],
domain="finn.custom_op.general")
modelproto = helper.make_model(
helper.make_graph([node_def], "test_model", [x], [out],
value_info=[W]))
model = ModelWrapper(modelproto)
model.set_tensor_datatype("x", DataType.BINARY)
model.set_tensor_datatype("W", DataType.BINARY)
W_data = np.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32)
model.set_initializer("W", W_data)
# test shape inference
model = model.transform(InferShapes())
assert model.get_tensor_shape("out") == [M, N]
# test datatype inference
assert model.get_tensor_datatype("out") is DataType.FLOAT32
model = model.transform(InferDataTypes())
assert model.get_tensor_datatype("out") is DataType.UINT32
# test execution
x_data = np.asarray([[1, 0, 0]], dtype=np.float32)
inp_dict = {"x": x_data}
out_dict = oxe.execute_onnx(model, inp_dict)
Wb = 2 * W_data - 1
xb = 2 * x_data - 1
rb = np.matmul(xb, Wb)
assert (2 * out_dict["out"] - K == rb).all()
def test_round_thresholds():
v = helper.make_tensor_value_info("v", TensorProto.FLOAT, [1, 4])
thresholds = helper.make_tensor_value_info("thresholds", TensorProto.FLOAT,
[4, 1])
out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4])
node_def = helper.make_node("MultiThreshold", ["v", "thresholds"], ["out"],
domain="finn")
graph_def = helper.make_graph([node_def], "test_model", [v, thresholds],
[out])
model_def = helper.make_model(graph_def)
model = ModelWrapper(model_def)
threshold_val = np.asarray([[-1.1], [0.7], [2.3], [5.1]], dtype=np.float32)
model.set_initializer("thresholds", threshold_val)
model.set_tensor_datatype("v", DataType.INT8)
inp_dict_f = {"v": np.floor(threshold_val).T}
inp_dict_n = {"v": np.round(threshold_val).T}
inp_dict_c = {"v": np.ceil(threshold_val).T}
orig_f = oxe.execute_onnx(model, inp_dict_f)["out"]
orig_n = oxe.execute_onnx(model, inp_dict_n)["out"]
orig_c = oxe.execute_onnx(model, inp_dict_c)["out"]
assert model.get_tensor_datatype("thresholds") == DataType.FLOAT32
new_model = model.transform(RoundAndClipThresholds())
# rounded up thresholds should have same dtype as input
assert new_model.get_tensor_datatype("thresholds") == DataType.INT8
new_f = oxe.execute_onnx(new_model, inp_dict_f)["out"]
new_n = oxe.execute_onnx(new_model, inp_dict_n)["out"]
new_c = oxe.execute_onnx(new_model, inp_dict_c)["out"]
assert np.isclose(orig_f, new_f, atol=1e-3).all()
assert np.isclose(orig_n, new_n, atol=1e-3).all()
assert np.isclose(orig_c, new_c, atol=1e-3).all()
def test_infer_datatypes():
raw_m = get_data("finn.data", "onnx/mnist-conv/model.onnx")
model = ModelWrapper(raw_m)
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveReadableTensorNames())
# this model has no DataType info, so add some DataType annotation
# to make things a bit more exciting
model.set_tensor_datatype("global_in", DataType["UINT8"])
# Conv with int weights + inputs will have int output datatype
model.set_tensor_datatype("Conv_0_param0", DataType["INT4"])
model = model.transform(InferDataTypes())
assert model.get_tensor_datatype("global_in") == DataType["UINT8"]
assert model.get_tensor_datatype("Conv_0_out0") == DataType["INT32"]
assert model.get_tensor_datatype("Relu_0_out0") == DataType["FLOAT32"]
assert model.get_tensor_datatype("global_out") == DataType["FLOAT32"]
def test_infer_datatypes():
lfc = get_test_model_trained("LFC", 1, 1)
bo.export_finn_onnx(lfc, (1, 1, 28, 28), export_onnx_path)
model = ModelWrapper(export_onnx_path)
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveReadableTensorNames())
model = model.transform(InferDataTypes())
assert model.get_tensor_datatype("MatMul_0_out0") == DataType.INT32
assert model.get_tensor_datatype("MatMul_1_out0") == DataType.INT32
assert model.get_tensor_datatype("MatMul_2_out0") == DataType.INT32
assert model.get_tensor_datatype("MatMul_3_out0") == DataType.INT32
assert model.get_tensor_datatype("Sign_0_out0") == DataType.BIPOLAR
assert model.get_tensor_datatype("Sign_1_out0") == DataType.BIPOLAR
assert model.get_tensor_datatype("Sign_2_out0") == DataType.BIPOLAR
assert model.get_tensor_datatype("Sign_3_out0") == DataType.BIPOLAR
os.remove(export_onnx_path)
def test_infer_datatypes():
raw_m = get_data("finn.data", "onnx/mnist-conv/model.onnx")
model = ModelWrapper(raw_m)
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveReadableTensorNames())
# this model has no DataType info, so add some DataType annotation
# to make things a bit more exciting
model.set_tensor_datatype("global_in", DataType.UINT8)
# manual non-float annotations on regular ONNX nodes won't disappear
# (InferDataTypes assumes they've been put there with good reason)
model.set_tensor_datatype("MaxPool_1_out0", DataType.INT4)
# MatMul with int weights + inputs will have int output datatype
model.set_tensor_datatype("MatMul_0_param0", DataType.UINT8)
model = model.transform(InferDataTypes())
assert model.get_tensor_datatype("global_in") == DataType.UINT8
assert model.get_tensor_datatype("Reshape_0_out0") == DataType.INT4
assert model.get_tensor_datatype("MatMul_0_out0") == DataType.INT32
assert model.get_tensor_datatype("global_out") == DataType.FLOAT32
def test_modelwrapper_setting_unsetting_datatypes():
# Set and unset some datatypes and check for expected return values
raw_m = get_data("finn.data", "onnx/mnist-conv/model.onnx")
model = ModelWrapper(raw_m)
test_node = model.graph.node[0]
test_tensor = test_node.output[0]
ret = model.get_tensor_datatype(test_tensor)
assert (ret == DataType["FLOAT32"]
), "Tensor datatype should be float32 for no initalization."
model.set_tensor_datatype(test_tensor, None)
ret = model.get_tensor_datatype(test_tensor)
assert ret == DataType[
"FLOAT32"], "An unset datatype should return float32."
model.set_tensor_datatype(test_tensor, DataType["INT3"])
ret = model.get_tensor_datatype(test_tensor)
assert ret == DataType["INT3"], "Tensor datatype should follow setting."
model.set_tensor_datatype(test_tensor, DataType["UINT4"])
ret = model.get_tensor_datatype(test_tensor)
assert ret == DataType["UINT4"], "Tensor datatype should follow setting."
model.set_tensor_datatype(test_tensor, None)
ret = model.get_tensor_datatype(test_tensor)
assert ret == DataType[
"FLOAT32"], "An unset datatype should return float32."
model.set_tensor_datatype(test_tensor, DataType["BIPOLAR"])
ret = model.get_tensor_datatype(test_tensor)
assert ret == DataType["BIPOLAR"], "Tensor datatype should follow setting."
def test_fpgadataflow_ipstitch_remote_execution():
try:
ip = os.environ["PYNQ_IP"] # NOQA
if ip == "":
pytest.skip("PYNQ board IP address not specified")
model = ModelWrapper(
ip_stitch_model_dir +
"/test_fpgadataflow_ipstitch_pynq_deployment.onnx")
iname = "inp"
idt = model.get_tensor_datatype(iname)
ishape = model.get_tensor_shape(iname)
x = gen_finn_dt_tensor(idt, ishape)
input_dict = {"inp": x}
outp = execute_onnx(model, input_dict)
assert np.isclose(outp["outp"], x).all()
except KeyError:
pytest.skip("PYNQ board IP address not specified")
def test_fpgadataflow_ipstitch_rtlsim():
model = ModelWrapper(ip_stitch_model_dir +
"/test_fpgadataflow_ip_stitch.onnx")
model.set_metadata_prop("rtlsim_trace", "whole_trace.vcd")
sim = pyverilate_stitched_ip(model)
exp_io = [
"ap_clk_0",
"ap_rst_n_0",
"in0_V_V_0_tdata",
"in0_V_V_0_tready",
"in0_V_V_0_tvalid",
"out_r_0_tdata",
"out_r_0_tkeep",
"out_r_0_tlast",
"out_r_0_tready",
"out_r_0_tvalid",
"s_axi_control_0_araddr",
"s_axi_control_0_arready",
"s_axi_control_0_arvalid",
"s_axi_control_0_awaddr",
"s_axi_control_0_awready",
"s_axi_control_0_awvalid",
"s_axi_control_0_bready",
"s_axi_control_0_bresp",
"s_axi_control_0_bvalid",
"s_axi_control_0_rdata",
"s_axi_control_0_rready",
"s_axi_control_0_rresp",
"s_axi_control_0_rvalid",
"s_axi_control_0_wdata",
"s_axi_control_0_wready",
"s_axi_control_0_wstrb",
"s_axi_control_0_wvalid",
]
assert dir(sim.io) == exp_io
model.set_metadata_prop("exec_mode", "rtlsim")
idt = model.get_tensor_datatype("inp")
ishape = model.get_tensor_shape("inp")
x = gen_finn_dt_tensor(idt, ishape)
# x = np.zeros(ishape, dtype=np.float32)
# x = np.asarray([[-2, -1, 0, 1]], dtype=np.float32)
rtlsim_res = execute_onnx(model, {"inp": x})["outp"]
assert (rtlsim_res == x).all()
def test_fpgadataflow_ipstitch_rtlsim():
model = ModelWrapper(ip_stitch_model_dir + "/test_fpgadataflow_ip_stitch.onnx")
sim = pyverilate_stitched_ip(model)
exp_io = [
"ap_clk_0",
"ap_rst_n_0",
"in0_V_V_0_tdata",
"in0_V_V_0_tready",
"in0_V_V_0_tvalid",
"out_r_0_tdata",
"out_r_0_tkeep",
"out_r_0_tlast",
"out_r_0_tready",
"out_r_0_tvalid",
]
assert dir(sim.io) == exp_io
model.set_metadata_prop("exec_mode", "rtlsim")
idt = model.get_tensor_datatype("inp")
ishape = model.get_tensor_shape("inp")
x = gen_finn_dt_tensor(idt, ishape)
rtlsim_res = execute_onnx(model, {"inp": x})["outp"]
assert (rtlsim_res == x).all()
def apply(self, model):
# create a temporary folder for the generated driver
pynq_driver_dir = make_build_dir(prefix="pynq_driver_")
model.set_metadata_prop("pynq_driver_dir", pynq_driver_dir)
# create the base FINN driver -- same for all accels
driver_base_template = pk.resource_filename(
"finn.qnn-data", "templates/driver/driver_base.py"
)
driver_base_py = pynq_driver_dir + "/driver_base.py"
shutil.copy(driver_base_template, driver_base_py)
# extract input-output shapes from the graph
# TODO convert this to an analysis pass?
idt = []
idma_names = []
ishape_normal = []
ishape_folded = []
ishape_packed = []
for idma_ind, graph_in in enumerate(model.graph.input):
i_tensor_name = graph_in.name
# get inp tensor properties
i_tensor_dt = model.get_tensor_datatype(i_tensor_name)
i_tensor_shape_normal = tuple(model.get_tensor_shape(i_tensor_name))
# go down into dataflow partition to get folded shape info etc
# TODO consider setting these as attributes during dataflow partitioning
i_consumer = model.find_consumer(i_tensor_name)
assert (
i_consumer.op_type == "StreamingDataflowPartition"
), """
Ensure CreateDataflowPartition called before driver creation."""
first_df_model = ModelWrapper(getCustomOp(i_consumer).get_nodeattr("model"))
assert (
first_df_model.graph.node[0].op_type == "IODMA"
), "First partition must hold input IODMA"
successors = model.find_direct_successors(i_consumer)
successor_input_num = list(successors[0].input).index(i_consumer.output[0])
successor_sdp = getCustomOp(successors[0])
successor_df_model = ModelWrapper(successor_sdp.get_nodeattr("model"))
first_node = successor_df_model.find_consumer(
successor_df_model.graph.input[successor_input_num].name
)
i_tensor_shape_folded = tuple(
getCustomOp(first_node).get_folded_input_shape()
)
# generate dummy folded i/o tensors and their packed versions
i_tensor_dummy_folded = gen_finn_dt_tensor(
i_tensor_dt, i_tensor_shape_folded
)
i_tensor_dummy_packed = dpk.finnpy_to_packed_bytearray(
i_tensor_dummy_folded, i_tensor_dt
)
i_tensor_shape_packed = i_tensor_dummy_packed.shape
# append all input tensor info to relevant lists
idt.append("DataType['%s']" % i_tensor_dt.name)
ishape_normal.append(i_tensor_shape_normal)
ishape_folded.append(i_tensor_shape_folded)
ishape_packed.append(i_tensor_shape_packed)
idma_names.append(getCustomOp(i_consumer).get_nodeattr("instance_name"))
odt = []
odma_names = []
oshape_normal = []
oshape_folded = []
oshape_packed = []
for odma_ind, graph_out in enumerate(model.graph.output):
o_tensor_name = graph_out.name
# get inp tensor properties
o_tensor_dt = model.get_tensor_datatype(o_tensor_name)
o_tensor_shape_normal = tuple(model.get_tensor_shape(o_tensor_name))
# go down into IODMA partition to get folded shape info etc
# TODO consider setting these as attributes during dataflow partitioning
o_producer = model.find_producer(o_tensor_name)
assert (
o_producer.op_type == "StreamingDataflowPartition"
), """
Ensure CreateDataflowPartition called before driver creation."""
df_model = ModelWrapper(getCustomOp(o_producer).get_nodeattr("model"))
assert (
df_model.graph.node[-1].op_type == "IODMA"
), "Partition must hold output IODMA"
predecessors = model.find_direct_predecessors(o_producer)
predecessor_output_num = list(predecessors[0].output).index(
o_producer.input[0]
)
predecessor_sdp = getCustomOp(predecessors[0])
predecessor_df_model = ModelWrapper(predecessor_sdp.get_nodeattr("model"))
last_node = predecessor_df_model.find_producer(
predecessor_df_model.graph.output[predecessor_output_num].name
)
o_tensor_shape_folded = tuple(
getCustomOp(last_node).get_folded_output_shape()
)
o_tensor_dummy_folded = gen_finn_dt_tensor(
o_tensor_dt, o_tensor_shape_folded
)
o_tensor_dummy_packed = dpk.finnpy_to_packed_bytearray(
o_tensor_dummy_folded, o_tensor_dt
)
o_tensor_shape_packed = o_tensor_dummy_packed.shape
# append all output tensor info to relevant lists
odt.append("DataType['%s']" % o_tensor_dt.name)
oshape_normal.append(o_tensor_shape_normal)
oshape_folded.append(o_tensor_shape_folded)
oshape_packed.append(o_tensor_shape_packed)
odma_names.append(getCustomOp(o_producer).get_nodeattr("instance_name"))
# generate external weights npy files
weights_dir = pynq_driver_dir + "/runtime_weights"
os.makedirs(weights_dir)
idma_idx = 0
ext_weight_dma_cnt = 0
for node in model.graph.node:
assert (
node.op_type == "StreamingDataflowPartition"
), "CreateDataflowPartition needs to be applied before driver generation"
if len(node.input) > 0:
producer = model.find_producer(node.input[0])
init_tensor = model.get_initializer(node.input[0])
else:
producer = None
init_tensor = None
if producer is None: # input dma?
sdp_inst = getCustomOp(node)
idma_name = sdp_inst.get_nodeattr("instance_name")
df_model = ModelWrapper(sdp_inst.get_nodeattr("model"))
assert df_model.graph.node[0].op_type == "IODMA"
iodma_node = getCustomOp(df_model.graph.node[0])
if iodma_node.get_nodeattr("burstMode") == "wrap": # input weights dma?
init_tensor = df_model.get_initializer(
iodma_node.onnx_node.input[0]
)
ext_weight_dma_cnt += 1
w_dtype = df_model.get_tensor_datatype(
iodma_node.onnx_node.input[0]
)
init_external_tensor = to_external_tensor(init_tensor, w_dtype)
np.save(
weights_dir + "/" + idma_name + ".npy", init_external_tensor
)
idma_idx += 1
# fill in the driver template
driver_py = pynq_driver_dir + "/driver.py"
driver = template_driver.pynq_driver_template
driver = driver.replace("$PLATFORM$", self.platform)
driver = driver.replace("$INPUT_FINN_DATATYPE$", str(idt).replace('"', ""))
driver = driver.replace("$INPUT_SHAPE_NORMAL$", str(ishape_normal))
driver = driver.replace("$INPUT_SHAPE_FOLDED$", str(ishape_folded))
driver = driver.replace("$INPUT_SHAPE_PACKED$", str(ishape_packed))
driver = driver.replace("$OUTPUT_FINN_DATATYPE$", str(odt).replace('"', ""))
driver = driver.replace("$OUTPUT_SHAPE_NORMAL$", str(oshape_normal))
driver = driver.replace("$OUTPUT_SHAPE_FOLDED$", str(oshape_folded))
driver = driver.replace("$OUTPUT_SHAPE_PACKED$", str(oshape_packed))
driver = driver.replace("$INPUT_DMA_NAME$", "%s" % str(idma_names))
driver = driver.replace("$OUTPUT_DMA_NAME$", "%s" % str(odma_names))
driver = driver.replace("$NUM_INPUTS$", str(len(idma_names)))
driver = driver.replace("$NUM_OUTPUTS$", str(len(odma_names)))
driver = driver.replace("$EXT_WEIGHT_NUM$", str(ext_weight_dma_cnt))
with open(driver_py, "w") as f:
f.write(driver)
# add validate.py to run full top-1 test (only for suitable networks)
validate_py = pynq_driver_dir + "/validate.py"
validate_template = pk.resource_filename(
"finn.qnn-data", "templates/driver/validate.py"
)
shutil.copy(validate_template, validate_py)
# copy all the dependencies into the driver folder
# driver imports utils/data_packing and core/datatype
# both of which are in finn-base
# e.g. /workspace/finn-base/src/finn/util/data_packing.py
dpk_root = dpk.__file__
# e.g. /workspace/finn-base/src/finn/util
dpk_root = dpk_root.replace("data_packing.py", "")
# e.g. /workspace/finn-base/src/finn/core/datatype.py
dtp_root = dtp.__file__
# e.g. /workspace/finn-base/src/finn/core
dtp_root = dtp_root.replace("datatype.py", "")
shutil.copytree(dpk_root, pynq_driver_dir + "/finn/util")
shutil.copytree(dtp_root, pynq_driver_dir + "/finn/core")
# generate weight files for runtime-writable layers
for sdp_ind, sdp_node in enumerate(model.graph.node):
assert sdp_node.op_type == "StreamingDataflowPartition"
# get dataflow model
sdp_node = getCustomOp(sdp_node)
dataflow_model_filename = sdp_node.get_nodeattr("model")
dataflow_model = ModelWrapper(dataflow_model_filename)
rt_layer_ind = 0
for node in dataflow_model.graph.node:
if node.op_type in ["StreamingFCLayer_Batch", "Thresholding_Batch"]:
node_inst = getCustomOp(node)
is_rt_weights = node_inst.get_nodeattr("runtime_writeable_weights")
if is_rt_weights == 1:
fcl_w = dataflow_model.get_initializer(node.input[1])
w_filename = weights_dir + "/%d_%d_%s.dat" % (
sdp_ind,
rt_layer_ind,
node.name,
)
node_inst.make_weight_file(
fcl_w, "decoupled_runtime", w_filename
)
rt_layer_ind += 1
elif node.op_type == "StreamingDataflowPartition":
warnings.warn(
"""Nested StreamingDataflowPartition are not supported
"""
)
else:
continue
return (model, False)
def execution_im2col(
x,
idt,
k_h,
k_w,
stride,
ifm_ch,
ifm_dim_h,
ifm_dim_w,
pad_amt,
pad_val=0,
dilation=1,
):
pad_amt_h = pad_amt[0] + pad_amt[2]
pad_amt_w = pad_amt[1] + pad_amt[3]
ofm_dim_h = compute_conv_output_dim(ifm_dim_h, k_h, stride, pad_amt_h, dilation)
ofm_dim_w = compute_conv_output_dim(ifm_dim_w, k_w, stride, pad_amt_w, dilation)
# set up onnx model
inp = helper.make_tensor_value_info(
"inp", TensorProto.FLOAT, [1, ifm_dim_h, ifm_dim_w, ifm_ch]
)
outp = helper.make_tensor_value_info(
"outp", TensorProto.FLOAT, [1, ofm_dim_h, ofm_dim_w, k_h * k_w * ifm_ch]
)
im2col_node = helper.make_node(
"Im2Col",
["inp"],
["outp"],
domain="finn.custom_op.general",
stride=stride,
kernel_size=[k_h, k_w],
pad_amount=pad_amt,
pad_value=pad_val,
input_shape="(1,{},{},{})".format(ifm_dim_h, ifm_dim_w, ifm_ch),
dilations=dilation,
)
graph = helper.make_graph(
nodes=[im2col_node], name="im2col_graph", inputs=[inp], outputs=[outp]
)
model = helper.make_model(graph, producer_name="im2col-model")
model = ModelWrapper(model)
model.set_tensor_datatype("inp", idt)
# test shape inference
model.transform(InferShapes())
assert model.get_tensor_shape("outp") == [
1,
ofm_dim_h,
ofm_dim_w,
k_h * k_w * ifm_ch,
]
# test datatype inference
assert model.get_tensor_datatype("outp") is DataType.FLOAT32
model = model.transform(InferDataTypes())
assert model.get_tensor_datatype("outp") is idt
# prepare input data
input_dict = {"inp": x}
# execute model
y_produced = oxe.execute_onnx(model, input_dict)["outp"]
return y_produced
def test_end2end_cybsec_mlp_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")
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)
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
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.graph.node[3].input[1]
assert finn_model.get_tensor_datatype(first_matmul_w_name) == DataType.INT2
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"]
