diff --git a/paddleslim/quant/quanter.py b/paddleslim/quant/quanter.py index 95ee7a049edea4040de83fdc4398b8d4a685e8ad..328861720ba92e5742fa4ae9e70bc0c329a7acb3 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)