def onnx2onnx_flow(m: onnx.ModelProto,
disable_fuse_bn=False,
bn_on_skip=False,
bn_before_add=False,
bgr=False,
norm=False,
rgba2yynn=False,
eliminate_tail=False) -> onnx.ModelProto:
"""Optimize the onnx.
Args:
m (ModelProto): the input onnx ModelProto
disable_fuse_bn (bool, optional): do not fuse BN into Conv. Defaults to False.
bn_on_skip (bool, optional): add BN operator on skip branches. Defaults to False.
bn_before_add (bool, optional): add BN before Add node on every branches. Defaults to False.
bgr (bool, optional): add an Conv layer to convert rgb input to bgr. Defaults to False.
norm (bool, optional): add an Conv layer to add 0.5 tp the input. Defaults to False.
rgba2yynn (bool, optional): add an Conv layer to convert rgb input to yynn . Defaults to False.
eliminate_tail (bool, optional): remove the trailing NPU unsupported nodes. Defaults to False.
Returns:
ModelProto: the optimized onnx model object.
"""
# temp.weight_broadcast(m.graph)
m = combo.preprocess(m, disable_fuse_bn)
# temp.fuse_bias_in_consecutive_1x1_conv(m.graph)
# Add BN on skip branch
if bn_on_skip:
other.add_bn_on_skip_branch(m.graph)
elif bn_before_add:
other.add_bn_before_add(m.graph)
other.add_bn_before_activation(m.graph)
# My optimization
m = combo.common_optimization(m)
# Special options
if bgr:
special.change_input_from_bgr_to_rgb(m)
if norm:
special.add_0_5_to_normalized_input(m)
if rgba2yynn:
special.add_rgb2yynn_node(m)
# Remove useless last node
if eliminate_tail:
eliminating.remove_useless_last_nodes(m.graph)
# Postprocessing
m = combo.postprocess(m)
return m
m = onnx.load(args.in_file)
# temp.weight_broadcast(m.graph)
m = combo.preprocess(m, args.disable_fuse_bn)
# temp.fuse_bias_in_consecutive_1x1_conv(m.graph)
# Add BN on skip branch
if args.bn_on_skip:
other.add_bn_on_skip_branch(m.graph)
elif args.bn_before_add:
other.add_bn_before_add(m.graph)
other.add_bn_before_activation(m.graph)
# Split deconv
if args.split_convtranspose:
other.split_ConvTranspose(m)
# My optimization
m = combo.common_optimization(m)
# Special options
if args.bgr:
special.change_input_from_bgr_to_rgb(m)
if args.norm:
special.add_0_5_to_normalized_input(m)
# Remove useless last node
if args.eliminate_tail:
eliminating.remove_useless_last_nodes(m.graph)
# Postprocessing
m = combo.postprocess(m)
onnx.save(m, outfile)