From 81db340ffffe9411523a273f52e51d74098010fc Mon Sep 17 00:00:00 2001 From: yukavio <67678385+yukavio@users.noreply.github.com> Date: Mon, 31 Aug 2020 11:41:49 +0800 Subject: [PATCH] Fix quant aware (#431) * fix pruned model save and load, fix bugs of pruning depthwise conv * add unit test of pruned model save and load * temp delete * add unit test * paddleslim/quant/quanter.py * fix bug --- paddleslim/quant/quanter.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/paddleslim/quant/quanter.py b/paddleslim/quant/quanter.py index 95ee7a04..32886172 100755 --- a/paddleslim/quant/quanter.py +++ b/paddleslim/quant/quanter.py @@ -185,7 +185,8 @@ def quant_aware(program, weight_preprocess_func=None, act_preprocess_func=None, optimizer_func=None, - executor=None): + executor=None, + return_program=False): """Add quantization and dequantization operators to "program" for quantization training or testing. @@ -226,6 +227,8 @@ def quant_aware(program, quantization function and preprocess function, this function must be set. Default is None. exe(Fluid.Executor): If user want to use self-defined quantization function and preprocess function, exe must be set for initialization. Default is None. + return_program(bool): If user want return value is a Program rather than Compiled Program, This argument should be set True. + Default is False. Returns: fluid.CompiledProgram | fluid.Program: Program with quantization and dequantization ``operators`` """ @@ -291,7 +294,7 @@ def quant_aware(program, VARS_MAPPING_TABLE)) save_dict(main_graph.out_node_mapping_table) - if for_test: + if for_test or return_program: quant_program = main_graph.to_program() else: quant_program = fluid.CompiledProgram(main_graph.graph) -- GitLab