def test_modelwrapper():
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)
assert model.check_all_tensor_shapes_specified() is False
inp_name = model.graph.input[0].name
inp_shape = model.get_tensor_shape(inp_name)
assert inp_shape == [1, 1, 28, 28]
# find first matmul node
l0_mat_tensor_name = ""
l0_inp_tensor_name = ""
for node in model.graph.node:
if node.op_type == "MatMul":
l0_inp_tensor_name = node.input[0]
l0_mat_tensor_name = node.input[1]
break
assert l0_mat_tensor_name != ""
l0_weights = model.get_initializer(l0_mat_tensor_name)
assert l0_weights.shape == (784, 1024)
l0_weights_hist = Counter(l0_weights.flatten())
assert (l0_weights_hist[1.0] + l0_weights_hist[-1.0]) == 784 * 1024
l0_weights_rand = np.random.randn(784, 1024)
model.set_initializer(l0_mat_tensor_name, l0_weights_rand)
assert (model.get_initializer(l0_mat_tensor_name) == l0_weights_rand).all()
assert l0_inp_tensor_name != ""
inp_cons = model.find_consumer(l0_inp_tensor_name)
assert inp_cons.op_type == "MatMul"
out_prod = model.find_producer(l0_inp_tensor_name)
assert out_prod.op_type == "MultiThreshold"
inp_layout = model.get_tensor_layout(inp_name)
assert inp_layout is None
inp_layout = DataLayout.NCHW
model.set_tensor_layout(inp_name, inp_layout)
assert model.get_tensor_layout(inp_name) == inp_layout
inp_sparsity = model.get_tensor_sparsity(inp_name)
assert inp_sparsity is None
inp_sparsity = {"dw": {"kernel_shape": 3}}
model.set_tensor_sparsity(inp_name, inp_sparsity)
assert model.get_tensor_sparsity(inp_name) == inp_sparsity
os.remove(export_onnx_path)
def test_modelwrapper():
raw_m = get_data("finn.data", "onnx/mnist-conv/model.onnx")
model = ModelWrapper(raw_m)
assert model.check_all_tensor_shapes_specified() is True
inp_name = model.graph.input[0].name
inp_shape = model.get_tensor_shape(inp_name)
assert inp_shape == [1, 1, 28, 28]
conv_nodes = model.get_nodes_by_op_type("Conv")
matmul_nodes = model.get_nodes_by_op_type("MatMul")
assert len(conv_nodes) == 2
assert len(matmul_nodes) == 1
first_conv = conv_nodes[0]
first_conv_iname = first_conv.input[0]
first_conv_wname = first_conv.input[1]
first_conv_oname = first_conv.output[0]
assert first_conv_iname != "" and (first_conv_iname is not None)
assert first_conv_wname != "" and (first_conv_wname is not None)
assert first_conv_oname != "" and (first_conv_oname is not None)
first_conv_weights = model.get_initializer(first_conv_wname)
assert first_conv_weights.shape == (8, 1, 5, 5)
first_conv_weights_rand = np.random.randn(8, 1, 5, 5)
model.set_initializer(first_conv_wname, first_conv_weights_rand)
assert (model.get_initializer(first_conv_wname) == first_conv_weights_rand
).all()
inp_cons = model.find_consumer(first_conv_iname)
assert inp_cons == first_conv
out_prod = model.find_producer(first_conv_oname)
assert out_prod == first_conv
inp_layout = model.get_tensor_layout(first_conv_iname)
assert inp_layout is None
inp_layout = DataLayout.NCHW
model.set_tensor_layout(first_conv_iname, inp_layout)
assert model.get_tensor_layout(first_conv_iname) == inp_layout
inp_sparsity = model.get_tensor_sparsity(first_conv_iname)
assert inp_sparsity is None
inp_sparsity = {"dw": {"kernel_shape": [3, 3]}}
model.set_tensor_sparsity(first_conv_iname, inp_sparsity)
assert model.get_tensor_sparsity(first_conv_iname) == inp_sparsity
def test_merge_onnx_models():
# load pre model
raw_m = get_data("finn", "data/onnx/mnist-conv/model.onnx")
model1 = ModelWrapper(raw_m)
# the input for model1 comes from a uint8 vector so we set the finn datatype
# of the input tensor to DataType.UINT8 to verify that the datatypes are correctly
# preserved in the transformed model
model1.set_tensor_datatype(model1.graph.input[0].name, DataType.UINT8)
model1 = model1.transform(InferShapes())
model1 = model1.transform(GiveUniqueNodeNames())
model1 = model1.transform(GiveReadableTensorNames())
# set up post model
shape = [1, 10]
inp = helper.make_tensor_value_info("inp", TensorProto.FLOAT, shape)
a0 = helper.make_tensor_value_info("a0", TensorProto.FLOAT, [])
a1 = helper.make_tensor_value_info("a1", TensorProto.FLOAT, [])
outp = helper.make_tensor_value_info("outp", TensorProto.FLOAT, shape)
mul_node = helper.make_node("Mul", ["inp", "a0"], ["mul_out"])
div_node = helper.make_node("Div", ["mul_out", "a1"], ["outp"])
graph = helper.make_graph(
nodes=[mul_node, div_node],
name="model2-graph",
inputs=[inp],
outputs=[outp],
value_info=[a0, a1],
)
model2 = helper.make_model(graph, producer_name="model2")
model2 = ModelWrapper(model2)
# initialize model2
a0_value = np.random.uniform(low=0, high=1, size=(1)).astype(np.float32)
model2.set_initializer("a0", a0_value)
a1_value = np.random.uniform(low=0.1, high=1, size=(1)).astype(np.float32)
model2.set_initializer("a1", a1_value)
# set a dummy sparsity annotation to check if it gets correctly transferred
# to the merged model
sparsity = {"dw": {"kernel_shape": 0}}
model2.set_tensor_sparsity("a1", sparsity)
model2 = model2.transform(InferShapes())
model2 = model2.transform(InferDataTypes())
model2 = model2.transform(InferDataLayouts())
model2 = model2.transform(GiveUniqueNodeNames())
model2 = model2.transform(GiveReadableTensorNames())
# simulate the models before the merging and pass the output of model1 to model2
# load one of the test vectors
raw_i = get_data("finn", "data/onnx/mnist-conv/test_data_set_0/input_0.pb")
inp_values = onnx.load_tensor_from_string(raw_i)
inp_values = np_helper.to_array(inp_values)
idict = {model1.graph.input[0].name: inp_values}
odict = oxe.execute_onnx(model1, idict)
temp = odict[model1.graph.output[0].name]
idict = {model2.graph.input[0].name: temp}
odict = oxe.execute_onnx(model2, idict)
outp = odict[model2.graph.output[0].name]
# merge models
model_transformed = model2.transform(MergeONNXModels(model1))
idict = {model_transformed.graph.input[0].name: inp_values}
odict = oxe.execute_onnx(model_transformed, idict)
outp_transformed = odict[model_transformed.graph.output[0].name]
assert (outp == outp_transformed).all()
assert len(model_transformed.graph.node) == len(model1.graph.node) + len(
model2.graph.node
)
# to test if the value is preserved we set the sparsity annotation of input[1]
# of the division block to a dummy value, we can now look for the division block
# and check if the sparsity annotation is still the same
for n in model_transformed.graph.node:
if n.op_type == "Div":
tensor_name = n.input[1]
set_sparsity = model_transformed.get_tensor_sparsity(tensor_name)
assert sparsity == set_sparsity
# check if finn datatype of graph.input[0] is still set to UINT8
assert model_transformed.get_tensor_datatype("global_in") == DataType.UINT8
def set_up_reference_model(act, idt, wdt, k, ifm_dim, ifm_ch, stride, padding):
# set up reference model consisting of Im2Col + MatMul (+ MultiThreshold)
ofm_ch = ifm_ch
ofm_dim = compute_conv_output_dim(ifm_dim, k, stride, pad=padding)
if act is None:
odt = DataType.INT32
else:
odt = act
out_act = oh.make_tensor_value_info("out_act", TensorProto.FLOAT,
[1, ofm_dim, ofm_dim, ofm_ch])
T = oh.make_tensor_value_info("T", TensorProto.FLOAT, [ofm_ch, 15])
tdt = DataType.INT32
thresh_node = oh.make_node(
"MultiThreshold",
domain="finn.custom_op.general",
inputs=["outp", "T"],
outputs=["out_act"],
data_layout="NHWC",
out_dtype=odt.name,
out_scale=1.0,
out_bias=0.0,
)
# set up onnx model
inp = oh.make_tensor_value_info("inp", TensorProto.FLOAT,
[1, ifm_dim, ifm_dim, ifm_ch])
outp = oh.make_tensor_value_info("outp", TensorProto.FLOAT,
[1, ofm_dim, ofm_dim, ofm_ch])
W_sparse = oh.make_tensor_value_info("W_sparse", TensorProto.FLOAT,
[ifm_ch * k * k, ofm_ch])
im2col_node = oh.make_node(
"Im2Col",
domain="finn.custom_op.general",
inputs=["inp"],
outputs=["im2col_out"],
kernel_size=k,
stride=stride,
pad_amount=padding,
input_shape="(1, {}, {}, {})".format(ifm_dim, ifm_dim, ifm_ch),
depthwise=1,
)
matmul_node = oh.make_node("MatMul",
inputs=["im2col_out", "W_sparse"],
outputs=["outp"])
if act is None:
node_list = [im2col_node, matmul_node]
global_out = outp
value_info = [W_sparse]
else:
node_list = [im2col_node, matmul_node, thresh_node]
global_out = out_act
value_info = [W_sparse, T]
graph = oh.make_graph(
nodes=node_list,
name="lowered_dw_cnv_graph",
inputs=[inp],
outputs=[global_out],
value_info=value_info,
)
model = oh.make_model(graph, producer_name="lowered_dw_cnv-model")
model = ModelWrapper(model)
# initialize model
model.set_tensor_datatype("inp", idt)
model.set_tensor_datatype(model.graph.output[0].name, odt)
model.set_tensor_datatype("W_sparse", wdt)
w_tensor = gen_finn_dt_tensor(wdt, [ofm_ch, 1, k, k])
# create sparse matrix
W_matrix = np.zeros((ofm_ch, ifm_ch, k, k))
for ch in range(ifm_ch):
W_matrix[ch][ch] = w_tensor[ch][0]
W_matrix = W_matrix.astype(np.float32)
W_matrix = W_matrix.transpose(0, 2, 3, 1)
W_matrix = W_matrix.reshape(ofm_ch, ifm_ch * k * k)
model.set_initializer("W_sparse", W_matrix.T)
sparsity = {"dw": {"kernel_shape": k}}
model.set_tensor_sparsity("W_sparse", sparsity)
if act is not None:
(min, max) = calculate_signed_dot_prod_range(idt, wdt, ifm_ch * k * k)
n_steps = odt.get_num_possible_values() - 1
T_values = np.random.randint(min, max - 1,
(ofm_ch, n_steps)).astype(np.float32)
# provide non-decreasing thresholds
T_values = np.sort(T_values, axis=1)
model.set_initializer("T", T_values)
model.set_tensor_datatype("T", tdt)
model = model.transform(InferShapes())
return model
