def test_end2end_mobilenet_export():
# export preprocessing
preproc_onnx = build_dir + "/end2end_mobilenet_preproc.onnx"
mean = [0.485, 0.456, 0.406]
std = 0.226
ch = 3
preproc = NormalizePreProc(mean, std, ch)
bo.export_finn_onnx(preproc, (1, 3, 224, 224), preproc_onnx)
preproc_model = ModelWrapper(preproc_onnx)
# set input finn datatype to UINT8
preproc_model.set_tensor_datatype(preproc_model.graph.input[0].name,
DataType["UINT8"])
preproc_model = preproc_model.transform(InferShapes())
preproc_model = preproc_model.transform(FoldConstants())
preproc_model = preproc_model.transform(GiveUniqueNodeNames())
preproc_model = preproc_model.transform(GiveUniqueParameterTensors())
preproc_model = preproc_model.transform(GiveReadableTensorNames())
preproc_model.save(build_dir + "/end2end_mobilenet_preproc.onnx")
# export mobilenet
finn_onnx = build_dir + "/end2end_mobilenet_export.onnx"
mobilenet = get_test_model_trained("mobilenet", 4, 4)
bo.export_finn_onnx(mobilenet, (1, 3, 224, 224), finn_onnx)
# calculate golden output with pytorch/brevitas and save as .npy
# get single image as input and prepare image
img = Image.open("/workspace/finn/tests/brevitas/king_charles.jpg")
# resize smallest side of the image to 256 pixels and resize larger side
# with same ratio
img = resize_smaller_side(256, img)
# crop central 224*224 window
img = crop_center(224, img)
# save image as numpy array and as torch tensor to enable testing in
# brevitas/pytorch and finn and transpose from (H, W, C) to (C, H, W)
img_np = np.asarray(img).copy().astype(np.float32).transpose(2, 0, 1)
img_np = img_np.reshape(1, 3, 224, 224)
np.save(build_dir + "/end2end_mobilenet_input.npy", img_np)
img_torch = torch.from_numpy(img_np).float()
# do forward pass in PyTorch/Brevitas
input_tensor = preproc.forward(img_torch)
golden = mobilenet.forward(input_tensor).detach().numpy()
golden_topk = golden.flatten()
golden_top5 = np.argsort(golden_topk)[-5:]
golden_top5 = np.flip(golden_top5)
golden_top5_prob = []
for index in golden_top5:
golden_top5_prob.append(golden_topk[index])
# save golden output values
np.save(build_dir + "/end2end_mobilenet_golden_top5.npy", golden_top5)
np.save(build_dir + "/end2end_mobilenet_golden_top5_prob.npy",
golden_top5_prob)
assert os.path.isfile(finn_onnx)
assert os.path.isfile(build_dir + "/end2end_mobilenet_preproc.onnx")
def inference_cost(model_filename,
*,
output_json=None,
output_onnx=None,
preprocess=True,
discount_sparsity=True):
"""Print the inference cost estimate metric for given ONNX model.
Supports the Quant op for weight/activation quantization.
:param model_filename: Filename for ONNX model
:param output_json: Optional JSON filename to save the inference cost dict
:param output_onnx: Optional ONNX filename to save the final model after any
preprocessing
:param preprocess: If set, run preprocessing steps such as shape inference,
datatype inference and constant folding. Strongly recommended.
:param discount_sparsity: If set, will discount op cost of MAC ops with a
constant zero weight, and the mem cost of constant zero weights.
"""
print("Inference cost for " + model_filename)
model = ModelWrapper(model_filename)
if preprocess:
qnt_nodes = model.get_nodes_by_op_type("Quant")
for qnt_node in qnt_nodes:
qnt_node.domain = "finn.custom_op.general"
model = model.transform(InferShapes())
model = model.transform(GiveUniqueParameterTensors())
model = model.transform(InferDataTypes())
model = model.transform(FoldConstants())
model = model.transform(RemoveUnusedTensors())
model = model.transform(RemoveStaticGraphInputs())
model = model.transform(InferDataTypes())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveReadableTensorNames())
if output_onnx is not None:
model.save(output_onnx)
ret = model.analysis(lambda x: infca.inference_cost(x, discount_sparsity))
bops = compute_bops(ret)
mem_w_bits = compute_mem_bits(ret, "mem_w")
mem_o_bits = compute_mem_bits(ret, "mem_o")
ret["total_bops"] = bops
ret["total_mem_w_bits"] = mem_w_bits
ret["total_mem_o_bits"] = mem_o_bits
if "unsupported" in ret:
ret["unsupported"] = str(ret["unsupported"])
print(json.dumps(ret, sort_keys=True, indent=2))
if output_json is not None:
with open(output_json, "w") as f:
json.dump(ret, f, sort_keys=True, indent=2)
def step_resnet50_tidy(model: ModelWrapper, cfg: DataflowBuildConfig):
model = model.transform(GiveUniqueParameterTensors())
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(RemoveStaticGraphInputs())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveReadableTensorNames())
model = model.transform(InferDataTypes())
model = model.transform(InsertTopK())
model = model.transform(InferShapes())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveReadableTensorNames())
model = model.transform(InferDataTypes())
return model
def test_end2end_mobilenet_tidy_and_merge_with_preproc():
preproc_model = load_test_checkpoint_or_skip(
build_dir + "/end2end_mobilenet_preproc.onnx")
model = load_test_checkpoint_or_skip(build_dir +
"/end2end_mobilenet_export.onnx")
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(InsertTopK())
# get initializer from Mul that will be absorbed into topk
a0 = model.get_initializer(model.graph.node[-2].input[1])
np.save(build_dir + "/end2end_mobilenet_topk_scale.npy", a0)
model = model.transform(absorb.AbsorbScalarMulAddIntoTopK())
model = model.transform(InferShapes())
model = model.transform(InferDataTypes())
model = model.transform(InferDataLayouts())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveUniqueParameterTensors())
model = model.transform(GiveReadableTensorNames())
model = model.transform(MergeONNXModels(preproc_model))
model.save(build_dir + "/end2end_mobilenet_tidy.onnx")
def test_give_unique_parameter_tensors():
# Create model
input_shape = [4, 4]
in1 = onnx.helper.make_tensor_value_info("in1", onnx.TensorProto.FLOAT, input_shape)
out1 = onnx.helper.make_tensor_value_info(
"out1", onnx.TensorProto.FLOAT, input_shape
)
graph_nodes = []
graph_nodes += [
onnx.helper.make_node("Add", inputs=["in1", "param1"], outputs=["t1"])
]
graph_nodes += [
onnx.helper.make_node("Sum", inputs=["t1", "param1", "param1"], outputs=["t2"])
]
graph_nodes += [
onnx.helper.make_node("Sum", inputs=["t2", "param2", "param1"], outputs=["t3"])
]
graph_nodes += [
onnx.helper.make_node("Add", inputs=["t3", "param1"], outputs=["out1"])
]
onnx_graph = onnx.helper.make_graph(
nodes=graph_nodes,
name="simple_graph",
inputs=[in1],
outputs=[out1],
)
onnx_model = onnx.helper.make_model(onnx_graph, producer_name="simple-model")
model = ModelWrapper(onnx_model)
# Set param values
np.random.seed(0)
param1 = np.random.rand(*input_shape).astype(np.float32)
param2 = np.random.rand(*input_shape).astype(np.float32)
model.set_initializer("param1", param1)
model.set_initializer("param2", param2)
model = model.transform(InferShapes())
# Apply transformation
new_model = model.transform(GiveUniqueParameterTensors())
new_model = new_model.transform(InferShapes())
# Test
# Breaks the model?
input_tensor = np.random.rand(*input_shape).astype(np.float32)
input_dict = {"in1": input_tensor}
# run original
expected_context = oxe.execute_onnx(model, input_dict)
expected_output = expected_context[model.graph.output[0].name]
# run modified
produced_context = oxe.execute_onnx(new_model, input_dict)
produced_output = produced_context[new_model.graph.output[0].name]
assert np.isclose(
expected_output, produced_output, atol=1e-8
).all(), " GiveUniqueParameterTensors() transform breaks the model"
# Does the job?
param_set = set()
param_cnt = 0
for n in new_model.graph.node:
for i in range(1, len(n.input)):
param_set |= {n.input[i]}
param_cnt += 1
assert len(param_set) == param_cnt, " There are still parameters reused"
def test_brevitas_compare_exported_mobilenet():
if "IMAGENET_VAL_PATH" not in os.environ.keys():
pytest.skip("Can't do validation without IMAGENET_VAL_PATH")
n_images = 10
debug_mode = False
export_onnx_path = make_build_dir("test_brevitas_mobilenet-v1_")
# export preprocessing
preproc_onnx = export_onnx_path + "/quant_mobilenet_v1_4b_preproc.onnx"
preproc = NormalizePreProc(mean, std, ch)
bo.export_finn_onnx(preproc, (1, 3, 224, 224), preproc_onnx)
preproc_model = ModelWrapper(preproc_onnx)
preproc_model = preproc_model.transform(InferShapes())
preproc_model = preproc_model.transform(GiveUniqueNodeNames())
preproc_model = preproc_model.transform(GiveUniqueParameterTensors())
preproc_model = preproc_model.transform(GiveReadableTensorNames())
# export the actual MobileNet-v1
finn_onnx = export_onnx_path + "/quant_mobilenet_v1_4b.onnx"
mobilenet = get_test_model_trained("mobilenet", 4, 4)
if debug_mode:
dbg_hook = bo.enable_debug(mobilenet)
bo.export_finn_onnx(mobilenet, (1, 3, 224, 224), finn_onnx)
model = ModelWrapper(finn_onnx)
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(RemoveStaticGraphInputs())
model = model.transform(InsertTopK())
# get initializer from Mul that will be absorbed into topk
a0 = model.get_initializer(model.get_nodes_by_op_type("Mul")[-1].input[1])
model = model.transform(absorb.AbsorbScalarMulAddIntoTopK())
model = model.transform(InferShapes())
model = model.transform(InferDataTypes())
model = model.transform(InferDataLayouts())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveUniqueParameterTensors())
model = model.transform(GiveReadableTensorNames())
model.save(export_onnx_path + "/quant_mobilenet_v1_4b_wo_preproc.onnx")
# create merged preprocessing + MobileNet-v1 model
model = model.transform(MergeONNXModels(preproc_model))
model.save(export_onnx_path + "/quant_mobilenet_v1_4b.onnx")
with open(
export_onnx_path + "/mobilenet_validation.csv", "w", newline=""
) as csvfile:
writer = csv.writer(csvfile)
writer.writerow(
[
"goldenID",
"brevitasTop5",
"brevitasTop5[%]",
"finnTop5",
"finnTop5[%]",
"top5equal",
"top5%equal",
]
)
csvfile.flush()
workload = imagenet_util.get_val_images(n_images, interleave_classes=True)
all_inds_ok = True
all_probs_ok = True
for (img_path, target_id) in workload:
img_np = imagenet_util.load_resize_crop(img_path)
img_torch = torch.from_numpy(img_np).float()
# do forward pass in PyTorch/Brevitas
input_tensor = preproc.forward(img_torch)
expected = mobilenet.forward(input_tensor).detach().numpy()
expected_topk = expected.flatten()
expected_top5 = np.argsort(expected_topk)[-5:]
expected_top5 = np.flip(expected_top5)
expected_top5_prob = []
for index in expected_top5:
expected_top5_prob.append(expected_topk[index])
idict = {model.graph.input[0].name: img_np}
odict = oxe.execute_onnx(model, idict, return_full_exec_context=True)
produced = odict[model.graph.output[0].name]
produced_prob = odict["TopK_0_out0"] * a0
inds_ok = (produced.flatten() == expected_top5).all()
probs_ok = np.isclose(produced_prob.flatten(), expected_top5_prob).all()
all_inds_ok = all_inds_ok and inds_ok
all_probs_ok = all_probs_ok and probs_ok
writer.writerow(
[
str(target_id),
str(expected_top5),
str(expected_top5_prob),
str(produced.flatten()),
str(produced_prob.flatten()),
str(inds_ok),
str(probs_ok),
]
)
csvfile.flush()
if ((not inds_ok) or (not probs_ok)) and debug_mode:
print("Results differ for %s" % img_path)
# check all tensors at debug markers
names_brevitas = set(dbg_hook.values.keys())
names_finn = set(odict.keys())
names_common = names_brevitas.intersection(names_finn)
for dbg_name in names_common:
if not np.isclose(
dbg_hook.values[dbg_name].detach().numpy(),
odict[dbg_name],
atol=1e-3,
).all():
print("Tensor %s differs between Brevitas and FINN" % dbg_name)
assert all_inds_ok and all_probs_ok
def apply(self, model):
graph_modified = False
pre_model = self.pre_model
post_model = copy.deepcopy(model)
# to avoid mix-ups, start by giving all tensors random names
pre_model = pre_model.transform(GiveRandomTensorNames())
post_model = post_model.transform(GiveRandomTensorNames())
# check for dynamic outputs of pre model
dyn_outp = []
for outp in pre_model.graph.output:
init_val = pre_model.get_initializer(outp.name)
if init_val is None:
dyn_outp.append(outp)
if len(dyn_outp) != 1:
warnings.warn(
"The pre model has more than one dynamic output! The transformation "
"tries to connect the first dynamic output to the first dynamic input "
"of the post model.")
# check for dynamic inputs of post model
dyn_inp = []
for inp in post_model.graph.input:
init_val = post_model.get_initializer(inp.name)
if init_val is None:
dyn_inp.append(inp)
if len(dyn_inp) != 1:
warnings.warn(
"The post model has more than one dynamic input! The transformation "
"tries to connect the first dynamic input to the first dynamic output "
"of the pre model.")
# erase all node names to avoid conflict
for n in pre_model.graph.node:
n.name = ""
for n in post_model.graph.node:
n.name = ""
# check if models can be merged
output_model_a = dyn_outp[0].name
input_model_b = dyn_inp[0].name
output_a_shape = pre_model.get_tensor_shape(output_model_a)
input_b_shape = post_model.get_tensor_shape(input_model_b)
assert (output_a_shape == input_b_shape
), "Models can't be merged! Shapes don't match."
# connect output of one model to input of the other
for n in pre_model.graph.node:
if output_model_a == n.output[0]:
n.output[0] = input_model_b
# extract information for new model
# nodes
node_pre = [node for node in pre_model.graph.node]
node_post = [node for node in post_model.graph.node]
node_new = node_pre + node_post
# in and output
inp = pre_model.graph.input[0]
outp = post_model.graph.output[0]
vi_pre = [x for x in pre_model.graph.value_info]
out_pre = [x for x in pre_model.graph.output]
qa_pre = [x for x in pre_model.graph.quantization_annotation]
init_pre = [x for x in pre_model.graph.initializer]
vi_post = [x for x in post_model.graph.value_info]
qa_post = [x for x in post_model.graph.quantization_annotation]
init_post = [x for x in post_model.graph.initializer]
vi_new = vi_pre + vi_post + out_pre
qa_new = qa_pre + qa_post
init_new = init_pre + init_post
# create new graph and model
new_graph = helper.make_graph(
nodes=node_new,
name="fuse-graph",
inputs=[inp],
outputs=[outp],
value_info=vi_new,
)
new_model = helper.make_model(new_graph, producer_name="fuse_model")
new_model = ModelWrapper(new_model)
for i in init_new:
new_model.graph.initializer.append(i)
for qa in qa_new:
new_model.graph.quantization_annotation.append(qa)
# tidy-up new model
model = new_model
model = model.transform(InferShapes())
model = model.transform(InferDataTypes())
model = model.transform(InferDataLayouts())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveUniqueParameterTensors())
model = model.transform(GiveReadableTensorNames())
return (model, graph_modified)
def test_brevitas_mobilenet():
# get single image as input and prepare image
img = Image.open("/workspace/finn/tests/brevitas/king_charles.jpg")
# resize smallest side of the image to 256 pixels and resize larger side
# with same ratio
img = resize_smaller_side(256, img)
# crop central 224*224 window
img = crop_center(224, img)
# save image as numpy array and as torch tensor to enable testing in
# brevitas/pytorch and finn and transpose from (H, W, C) to (C, H, W)
img_np = np.asarray(img).copy().astype(np.float32).transpose(2, 0, 1)
img_np = img_np.reshape(1, 3, 224, 224)
img_torch = torch.from_numpy(img_np).float()
# export preprocess
export_onnx_path = make_build_dir("test_brevitas_mobilenet-v1_")
preproc_onnx = export_onnx_path + "/quant_mobilenet_v1_4b_preproc.onnx"
mean = [0.485, 0.456, 0.406]
std = 0.226
ch = 3
preproc = NormalizePreProc(mean, std, ch)
bo.export_finn_onnx(preproc, (1, 3, 224, 224), preproc_onnx)
preproc_model = ModelWrapper(preproc_onnx)
# set input finn datatype to UINT8
preproc_model.set_tensor_datatype(preproc_model.graph.input[0].name, DataType.UINT8)
preproc_model = preproc_model.transform(InferShapes())
preproc_model = preproc_model.transform(GiveUniqueNodeNames())
preproc_model = preproc_model.transform(GiveUniqueParameterTensors())
preproc_model = preproc_model.transform(GiveReadableTensorNames())
finn_onnx = export_onnx_path + "/quant_mobilenet_v1_4b_exported.onnx"
mobilenet = get_test_model_trained("mobilenet", 4, 4)
bo.export_finn_onnx(mobilenet, (1, 3, 224, 224), finn_onnx)
# do forward pass in PyTorch/Brevitas
input_tensor = preproc.forward(img_torch)
expected = mobilenet.forward(input_tensor).detach().numpy()
expected_topk = expected.flatten()
expected_top5 = np.argsort(expected_topk)[-5:]
expected_top5 = np.flip(expected_top5)
expected_top5_prob = []
for index in expected_top5:
expected_top5_prob.append(expected_topk[index])
model = ModelWrapper(finn_onnx)
model = model.transform(InferShapes())
model = model.transform(FoldConstants())
model = model.transform(InsertTopK())
# get initializer from Mul that will be absorbed into topk
a0 = model.get_initializer(model.graph.node[-2].input[1])
model = model.transform(absorb.AbsorbScalarMulAddIntoTopK())
model = model.transform(InferShapes())
model = model.transform(InferDataTypes())
model = model.transform(InferDataLayouts())
model = model.transform(GiveUniqueNodeNames())
model = model.transform(GiveUniqueParameterTensors())
model = model.transform(GiveReadableTensorNames())
model.save(export_onnx_path + "/quant_mobilenet_v1_4b_wo_preproc.onnx")
model = model.transform(MergeONNXModels(preproc_model))
model.save(export_onnx_path + "/quant_mobilenet_v1_4b.onnx")
idict = {model.graph.input[0].name: img_np}
odict = oxe.execute_onnx(model, idict, True)
produced = odict[model.graph.output[0].name]
produced_prob = odict["TopK_0_out0"] * a0
assert (produced.flatten() == expected_top5).all()
assert np.isclose(produced_prob.flatten(), expected_top5_prob).all()
