Convert the test program

Author: juncaipeng

python/paddle/fluid/contrib/slim/quantization/quantize_transpiler_v2.py

@@ -28,6 +28,27 @@
     __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s')
 
 
+def find_next_ops(block, var_name):
+    """
+    Find all followed ops for the input variable.
+    """
+    res_ops = []
+    for op in block.ops:
+        if var_name in op.input_arg_names:
+            res_ops.append(op)
+    return res_ops
+
+
+def load_variable_data(scope, var_name):
+    '''
+    Load variable value from scope
+    '''
+    var_node = scope.find_var(var_name)
+    assert var_node is not None, \
+        "Cannot find " + var_name + " in scope."
+    return np.array(var_node.get_tensor())
+
+
 class QuantizeTranspilerV2(object):
     def __init__(self,
                  weight_bits=8,
@@ -118,25 +139,36 @@ def apply(self, program, startup_program, is_test=False):
 
     def convert(self, test_program, scope=None):
         """
-        Convert the test program.
+        Convert the test program. 
+        Get the out scale from the moving_average_abs_max_scale op and save the
+        out scale into the quantized op. 
         Args:
             test_program(Program): the test program to be converted.
             scope(fluid.Scope, optional): The scope of the program, use it to load 
                 and save variables. If scope=None, get scope by global_scope(). 
         """
         scope = global_scope() if scope == None else scope
 
-        target_ops = []
         for block in test_program.blocks:
             for op in block.ops:
-                if op.type == "moving_average_abs_max_scale":
-                    target_ops.append(op)
+                if op.has_attr("quantization_type") \
+                    and op.attr("quantization_type") == "qat_with_weight":
+                    # quant op -> var1 -> fake op -> var2
+                    assert len(op.output_arg_names) == 1
+                    var1_name = op.output_arg_names[0]
+
+                    fake_ops = find_next_ops(block, var1_name)
+                    assert len(fake_ops) == 1
+                    fake_op = fake_ops[0]
+                    assert fake_op.type == "moving_average_abs_max_scale"
+
+                    out_scale_name = fake_op.output("OutScale")
+                    out_threshold = load_variable_data(scope, out_scale_name[0])
+                    op._set_attr("out_threshold", float(out_threshold))
 
-        for op in target_ops:
-            out_scale_name = op.output("OutScale")
-            # TODO: save the out threshold in the target ops
-            #print(out_scale_name)
-            #print(self._load_variable_data(scope, out_scale_name[0]))
+                    var2_name = fake_op.output("Out")[0]
+                    op._rename_output(var1_name, var2_name)
+                    fake_op._rename_output(var2_name, var1_name)
 
     def _transform_forward(self, block, op, var_rename_map, is_test):
         """
@@ -183,7 +215,7 @@ def _transform_forward(self, block, op, var_rename_map, is_test):
 
         # insert out scale op followed the quantized op
         for out_name in op.output_arg_names:
-            next_ops = self._find_next_ops(block, out_name)
+            next_ops = find_next_ops(block, out_name)
 
             idx = block.ops.index(op)
             out_var = block.var(out_name)
@@ -194,25 +226,6 @@ def _transform_forward(self, block, op, var_rename_map, is_test):
                 if "_grad" not in next_op.type:
                     next_op._rename_input(out_name, new_out_var.name)
 
-    def _find_next_ops(self, block, var_name):
-        """
-        Find all followed ops for the input variable.
-        """
-        res_ops = []
-        for op in block.ops:
-            if var_name in op.input_arg_names:
-                res_ops.append(op)
-        return res_ops
-
-    def _load_variable_data(self, scope, var_name):
-        '''
-        Load variable value from scope
-        '''
-        var_node = scope.find_var(var_name)
-        assert var_node is not None, \
-            "Cannot find " + var_name + " in scope."
-        return np.array(var_node.get_tensor())
-
     def _is_skip_quant(self, op):
         """
         Analyse whether the op should skip quantization or not.

python/paddle/fluid/contrib/slim/tests/test_quantize_transpiler_v2.py

@@ -79,6 +79,7 @@ def build_program(main, startup, is_test):
         random.seed(0)
         np.random.seed(0)
 
+        # 1 Define program
         train_program = fluid.Program()
         startup_program = fluid.Program()
         test_program = fluid.Program()
@@ -93,13 +94,14 @@ def build_program(main, startup, is_test):
             test_graph = IrGraph(core.Graph(test_program.desc), for_test=True)
             test_graph.draw('.', 'test_program_1')
 
+        # 2 Apply quantization
         qt = QuantizeTranspilerV2(
             activation_quantize_type=activation_quant_type,
-            weight_quantize_type=weight_quant_type,
-            quantizable_op_type=['conv2d', 'depthwise_conv2d', 'mul'])
-        qt.apply(train_program, startup_program, False)
-        qt.apply(test_program, startup_program, True)
+            weight_quantize_type=weight_quant_type)
+        qt.apply(train_program, startup_program, is_test=False)
+        qt.apply(test_program, startup_program, is_test=True)
 
+        # 3 Train
         place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
         exe = fluid.Executor(place)
         scope = fluid.Scope()
@@ -135,6 +137,8 @@ def build_program(main, startup, is_test):
 
         print('{}: {}'.format('loss', np.mean(loss_v)))
 
+        # 4 Convert
+        qt.convert(test_program, scope)
         if not for_ci:
             with fluid.scope_guard(scope):
                 fluid.io.save_inference_model('./infer_model',