From 0d9b45e8042ee3a7f4469a14c3a708eb407e757f Mon Sep 17 00:00:00 2001 From: sunnycase Date: Sat, 12 Feb 2022 00:19:31 +0800 Subject: [PATCH] Allow cut postprocess but preserve input (#1275) * Allow cut postprocess but preserve input * apply code-format changes Co-authored-by: sunnycase --- examples/tm_mobilenet_ssd_uint8_timvx.cpp | 4 ++-- tools/optimize/yolov5s-opt.py | 14 ++++++++------ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/examples/tm_mobilenet_ssd_uint8_timvx.cpp b/examples/tm_mobilenet_ssd_uint8_timvx.cpp index ec86b4dc..14a1888f 100644 --- a/examples/tm_mobilenet_ssd_uint8_timvx.cpp +++ b/examples/tm_mobilenet_ssd_uint8_timvx.cpp @@ -185,9 +185,9 @@ int main(int argc, char* argv[]) // create VeriSilicon TIM-VX backend context_t timvx_context = create_context("timvx", 1); - if ( set_context_device(timvx_context, "TIMVX", nullptr, 0) < 0 ) + if (set_context_device(timvx_context, "TIMVX", nullptr, 0) < 0) { - fprintf(stderr, "add_context_device failed.\n" ); + fprintf(stderr, "add_context_device failed.\n"); return 1; } diff --git a/tools/optimize/yolov5s-opt.py b/tools/optimize/yolov5s-opt.py index 3e48ddad..74f40cec 100644 --- a/tools/optimize/yolov5s-opt.py +++ b/tools/optimize/yolov5s-opt.py @@ -60,7 +60,7 @@ def parse_args(): args = parse_args() -def cut_focus_output(input_node, in_name, out_name): +def cut_focus_output(input_node, in_name, out_name, cut_focus): """ cut the focus and postprocess nodes Args: @@ -98,8 +98,11 @@ def cut_focus_output(input_node, in_name, out_name): del input_node[i] # cut input node - for n in in_name: - new_nodes = input_node[(node_dict[n] + 1):] + if cut_focus: + for n in in_name: + new_nodes = input_node[(node_dict[n] + 1):] + else: + new_nodes = input_node[:] return new_nodes @@ -181,9 +184,8 @@ def main(): new_nodes = old_node[:] # cut the focus and postprocess nodes - if args.cut_focus: - print("[Quant Tools Info]: Step 1, Remove the focus and postprocess nodes.") - new_nodes = cut_focus_output(old_node, in_tensor, out_tensor) + print("[Quant Tools Info]: Step 1, Remove the focus and postprocess nodes.") + new_nodes = cut_focus_output(old_node, in_tensor, out_tensor, args.cut_focus) # op fusion, using HardSwish replace the Sigmoid and Mul print("[Quant Tools Info]: Step 2, Using hardswish replace the sigmoid and mul.") -- GitLab