# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import collections import numpy as np from paddle.fluid.framework import ( default_main_program, default_startup_program, program_guard, ) from paddle.fluid.layer_helper import LayerHelper from paddle.fluid import unique_name from paddle.fluid import core from paddle.fluid.initializer import Constant from paddle.fluid.param_attr import ParamAttr from paddle.fluid.layer_helper import LayerHelper from paddle.fluid.layers.nn import autoincreased_step_counter from paddle.fluid.framework import Variable from paddle.fluid.executor import global_scope __all__ = ['QuantizeTranspiler'] _QUANTIZABLE_OP_TYPES = ['conv2d', 'depthwise_conv2d', 'mul'] def _quantized_var_name(var_name): """ Return quantized variable name for the input `var_name`. """ return "%s.quantized" % (var_name) def _dequantized_var_name(var_name): """ Return dequantized variable name for the input `var_name`. """ return "%s.dequantized" % (var_name) def _quantized_scale_name(var_name): """ Return quantized variable name for the input `var_name`. """ return "%s.scale" % (var_name) def _original_var_name(var_name): """ Return the original variable name. """ if var_name.endswith('.quantized.dequantized'): return var_name[: -len('.quantized.dequantized')] if var_name.endswith('.quantized'): return var_name[: -len('.quantized')] if var_name.endswith('.dequantized'): return var_name[: -len('.dequantized')] if var_name.endswith('.scale'): return var_name[: -len('.scale')] else: return var_name def _is_float(v): return isinstance(v, float) or isinstance(v, np.float32) def quant(x, scale, num_bits): y = np.round(x / scale * ((1 << (num_bits - 1)) - 1)) return y class QuantizeTranspiler: def __init__( self, weight_bits=8, activation_bits=8, activation_quantize_type='abs_max', weight_quantize_type='abs_max', window_size=10000, moving_rate=0.9, ): """ Convert and rewrite the fluid Program according to weight and activation quantization type. Args: weight_bits (int): quantization bit number for weights, the bias is not quantized. activation_bits (int): quantization bit number for activation. activation_quantize_type (str): quantization type for activation, now support 'abs_max', 'range_abs_max'. If use 'abs_max' mode, the quantization scale will be calculated dynamically each step in both training and testing period. If use 'range_abs_max', a static quantization scale will be calculated during training and used in inference. weight_quantize_type (str): quantization type for weights, support 'abs_max'. The 'range_abs_max' usually is not used for weight, since weights are fixed once the model is well trained. window_size (int): the window size for 'range_abs_max' quantization. Examples: .. code-block:: python # the original program will be rewrite, if you don't want to # change it, please clone at first. # quantize_program = program.clone() t = fluid.QuantizeTranspiler() t.transpile(quantize_program) """ self.weight_bits = weight_bits self.activation_bits = activation_bits quant_type = ['abs_max', 'range_abs_max', 'moving_average_abs_max'] if weight_quantize_type not in quant_type: raise ValueError( "Unknown weight_quantize_type: '%s'. It can only be ", "'abs_max' or 'range_abs_max' or 'moving_average_abs_max'.", str(weight_quantize_type), ) if activation_quantize_type not in quant_type: raise ValueError( "Unknown activation_quantize_type : '%s'. It can only be ", "'abs_max' or 'range_abs_max' or 'moving_average_abs_max'.", str(activation_quantize_type), ) self.weight_quantize_type = weight_quantize_type self.activation_quantize_type = activation_quantize_type self.window_size = window_size self.moving_rate = moving_rate self.helper = LayerHelper(self.__class__.__name__) self.fake_quant_op_types = [ 'fake_quantize_abs_max', 'fake_quantize_range_abs_max', 'fake_quantize_moving_average_abs_max', ] self.fake_dequant_op_types = ['fake_dequantize_max_abs'] self.is_test = None self.global_step = None def training_transpile(self, program=None, startup_program=None): """Rewrites a training input program in place for simulated quantization. Insert fake quantization and de-quantization ops into program to simulate the error introduced by quantization. And change the gradient ops' input by using the faked quantization weights and activation. Since the program is transformed in place, the graph connection will change. Args: program (Program): the input program to be transpile. """ self.is_test = False program = default_main_program() if program is None else program startup_program = ( default_startup_program() if startup_program is None else startup_program ) # marked the variable which has been quantized and dequantized. dequanted_vars = [ collections.OrderedDict() for _ in range(len(program.blocks)) ] grad_op_types = ['%s_grad' % (type) for type in _QUANTIZABLE_OP_TYPES] params = [p.name for p in program.global_block().iter_parameters()] def _transpile_forward(block, op): idx = block.ops.index(op) block_id = block.idx # insert quant op and dequant op for name in op.input_arg_names: # if share input between ops if name in dequanted_vars[block_id]: dequant_var = dequanted_vars[block_id][name] else: var = block.var(name) quant_bits = ( self.weight_bits if var.name in params else self.activation_bits ) quant_type = ( self.weight_quantize_type if var.name in params else self.activation_quantize_type ) quant_var, scale_var = self._insert_quant_op( block, idx, var, quant_bits, quant_type ) dequant_var = self._insert_dequant_op( block, idx + 1, quant_var, scale_var, quant_bits ) dequanted_vars[block_id][name] = dequant_var # rename the forward op inputs op._rename_input(name, dequant_var.name) def _transpile_backward(block, op): block_id = block.idx no_dequanted_input_vars = True for name in op.input_arg_names: if name in dequanted_vars[block_id]: dequant_var = dequanted_vars[block_id][name] op._rename_input(name, dequant_var.name) no_dequanted_input_vars = False if no_dequanted_input_vars: raise ValueError( "There is no dequanted inputs for op %s." % (op.type) ) with program_guard(program, startup_program): self._create_global_step() for block in program.blocks: ops = list(block.ops) block_id = block.idx for op in ops: # rewrite the forward ProgramDes if op.type in _QUANTIZABLE_OP_TYPES: _transpile_forward(block, op) # rename the backward op inputs if op.type in grad_op_types: _transpile_backward(block, op) def _create_global_step(self): if ( self.weight_quantize_type == 'range_abs_max' or self.activation_quantize_type == 'range_abs_max' ): self.global_step = autoincreased_step_counter() def freeze_program(self, program, place, scope=None): """Freeze input training program for inference. Args: program (Program): the input program to be transpile. """ self.is_test = True scope = global_scope() if scope is None else scope program = default_main_program() if program is None else program persistable_vars = [ v.name for v in filter(lambda var: var.persistable, program.list_vars()) ] op_in_rename_map = [ collections.OrderedDict() for _ in range(len(program.blocks)) ] op_out_rename_map = [ collections.OrderedDict() for _ in range(len(program.blocks)) ] var_scale_map = [ collections.OrderedDict() for _ in range(len(program.blocks)) ] def _remove_fake_quant_and_dequant_op(block, op): idx = block.ops.index(op) block_id = block.idx k = op.output('Out')[0] v = op.input('X')[0] if v not in op_in_rename_map[block_id]: op_in_rename_map[block_id][k] = v else: op_in_rename_map[block_id][k] = op_in_rename_map[block_id][v] block._remove_op(idx) def _insert_post_dequant_op(block, op): idx = block.ops.index(op) block_id = block.idx max_range = None scale_var = None for name in op.input_arg_names: # rename input name of the op to the input name of last op which has be removed if name in op_in_rename_map[block_id]: op._rename_input(name, op_in_rename_map[block_id][name]) scale_v = var_scale_map[block_id][_original_var_name(name)] if _original_var_name(name) in persistable_vars: param_range = (1 << (self.weight_bits - 1)) - 1 act_range = (1 << (self.activation_bits - 1)) - 1 assert _is_float(scale_v) max_range = param_range * act_range / scale_v else: assert isinstance(scale_v, Variable) scale_var = scale_v if len(op.output_arg_names) != 1: raise ValueError( "Only support one output, but op %s has" " more than one output." % (op.type) ) out_var = block.var(op.output_arg_names[0]) dequant_var = block.create_var( name=_dequantized_var_name(out_var.name), type=out_var.type, shape=out_var.shape, dtype=out_var.dtype, ) # insert fake_dequantize_op dequant_op = block._insert_op( idx + 1, type="fake_dequantize_max_abs", attrs={'max_range': float(max_range)}, inputs={"X": out_var, 'Scale': scale_var}, outputs={"Out": dequant_var}, ) op_out_rename_map[block_id][out_var.name] = dequant_var.name return dequant_var def _load_var(name): return np.array(scope.find_var(name).get_tensor()) def _restore_var(name, arr): t = scope.find_var(name).get_tensor() t.set(arr, place) for block in program.blocks: ops = list(block.ops) block_id = block.idx for op in ops: op_type = op.type # insert dequant_op after fc/conv, need to rename # input of the followed ops(of fc/conv) to the dquant_op for name in op.input_arg_names: if name in op_out_rename_map[block_id]: op._rename_input( name, op_out_rename_map[block_id][name] ) if op_type in self.fake_quant_op_types: in_arg_name = op.input('X')[0] if in_arg_name in persistable_vars: if self.weight_quantize_type == 'abs_max': param = _load_var(in_arg_name) scale_v = np.max(np.abs(param)) else: scale_v = _load_var(op.output('OutScale')[0]) var_scale_map[block_id][in_arg_name] = scale_v else: scale_v = block.var(op.output('OutScale')[0]) var_scale_map[block_id][in_arg_name] = scale_v if in_arg_name in persistable_vars: _remove_fake_quant_and_dequant_op(block, op) # quantize weight and restore param_t = _load_var(in_arg_name) param_q_t = quant(param_t, scale_v, self.weight_bits) _restore_var(in_arg_name, param_q_t) if op_type in self.fake_dequant_op_types: _remove_fake_quant_and_dequant_op(block, op) if op_type in _QUANTIZABLE_OP_TYPES: dequant_var = _insert_post_dequant_op(block, op) # remove the unused var in ProgramDesc self._remove_unused_var(program) # program = program.clone() def convert_to_int8(self, program, place, scope=None): scope = global_scope() if scope is None else scope program = default_main_program() if program is None else program def _load_var(name): return np.array(scope.find_var(name).get_tensor()) global_block = program.global_block() def convert_to_int8(var): int8_var_name = var.name + ".int8" int8_var = global_block.create_parameter( name=int8_var_name.encode('ascii'), type=var.type, dtype=core.VarDesc.VarType.INT8, shape=var.shape, ) tensor = _load_var(var.name) scope.var(int8_var_name) int8_tensor = scope.find_var(int8_var_name).get_tensor() int8_tensor.set(tensor.astype(np.int8), place) return int8_var input_map = {} for block in program.blocks: for op in list(block.ops): if op.type in _QUANTIZABLE_OP_TYPES: for name in op.input_arg_names: var = block.var(name) if var.persistable: if name not in input_map: int8_var = convert_to_int8(var) input_map[name] = int8_var.name op._rename_input(name, input_map[name]) self._remove_unused_var(program) def _remove_unused_var(self, program): all_remove_vars = [] for block in program.blocks: args = [] for op in block.ops: args += op.input_arg_names args += op.output_arg_names args = list(set(args)) # vals of all left ops var_names = block.vars.keys() # all vals sub_block_remove_vars = [] for var in var_names: if var not in args: sub_block_remove_vars.append(var) all_remove_vars.append(sub_block_remove_vars) remove_vars = [list(set(v)) for v in all_remove_vars] for i, block in enumerate(program.blocks): for v in remove_vars[i]: block._remove_var(v) def _insert_quant_abs_max_op(self, block, idx, var, quant_bits): """Insert fake_quantize_abs_max op.""" quant_var = block.create_var( name=_quantized_var_name(var.name), type=var.type, shape=var.shape, dtype=var.dtype, ) scale = block.create_var( name=_quantized_scale_name(var.name), type=var.type, shape=var.shape, dtype=var.dtype, ) quant_op = block._insert_op( idx, type='fake_quantize_abs_max', attrs={'bit_length': quant_bits}, inputs={'X': var}, outputs={'Out': quant_var, 'OutScale': scale}, ) return quant_var, scale def _insert_quant_range_abs_max_op(self, block, idx, var, quant_bits): """Insert fake_quantize_range_abs_max""" quant_var = block.create_var( name=_quantized_var_name(var.name), type=var.type, shape=var.shape, dtype=var.dtype, ) scale = self.helper.create_parameter( attr=ParamAttr( name=_quantized_scale_name(var.name), initializer=Constant(0.001), trainable=False, ), shape=[1], dtype=var.dtype, ) scale.stop_gradient = True ins = {'X': var, 'InScale': scale} outs = {'Out': quant_var, 'OutScale': scale} if not self.is_test: # A global step counter variable with type int64 scales = self.helper.create_global_variable( name=unique_name.generate('scales'), persistable=True, dtype=var.dtype, shape=[self.window_size], ) self.helper.set_variable_initializer( scales, initializer=Constant(value=0) ) ins['Iter'] = self.global_step outs['OutScales'] = scales attrs = { 'window_size': self.window_size, 'bit_length': quant_bits, 'is_test': self.is_test, } quant_op = block._insert_op( idx, type='fake_quantize_range_abs_max', attrs=attrs, inputs=ins, outputs=outs, ) return quant_var, scale def _insert_quant_moving_average_abs_max_op( self, block, idx, var, quant_bits ): """Insert fake_quantize_moving_average_abs_max""" quant_var = block.create_var( name=_quantized_var_name(var.name), type=var.type, shape=var.shape, dtype=var.dtype, ) state = self.helper.create_global_variable( name=unique_name.generate('state'), persistable=True, dtype=var.dtype, shape=[1], ) self.helper.set_variable_initializer( state, initializer=Constant(value=1) ) accum = self.helper.create_global_variable( name=unique_name.generate('accum'), persistable=True, dtype=var.dtype, shape=[1], ) self.helper.set_variable_initializer( accum, initializer=Constant(value=1) ) scale = self.helper.create_parameter( attr=ParamAttr( name=_quantized_scale_name(var.name), initializer=Constant(0.001), trainable=False, ), shape=[1], dtype=var.dtype, ) scale.stop_gradient = True ins = {'X': var, 'InScale': scale} outs = {'Out': quant_var, 'OutScale': scale} if not self.is_test: ins['InState'] = state ins['InAccum'] = accum outs['OutState'] = state outs['OutAccum'] = accum attrs = { 'bit_length': quant_bits, 'moving_rate': self.moving_rate, 'is_test': self.is_test, } quant_op = block._insert_op( idx, type='fake_quantize_moving_average_abs_max', attrs=attrs, inputs=ins, outputs=outs, ) return quant_var, scale def _insert_quant_op(self, block, idx, var, quant_bits, quant_type): """ Insert fake_quantize_op """ if quant_type == 'abs_max': return self._insert_quant_abs_max_op(block, idx, var, quant_bits) elif quant_type == 'range_abs_max': return self._insert_quant_range_abs_max_op( block, idx, var, quant_bits ) elif quant_type == 'moving_average_abs_max': return self._insert_quant_moving_average_abs_max_op( block, idx, var, quant_bits ) def _insert_dequant_op(self, block, idx, var, scale, quant_bits): """ Insert fake_quantize_op """ dequant_var = block.create_var( name=_dequantized_var_name(var.name), type=var.type, shape=var.shape, dtype=var.dtype, ) # insert fake_dequantize_op max_range = (1 << (quant_bits - 1)) - 1 dequant_op = block._insert_op( idx, type="fake_dequantize_max_abs", attrs={'max_range': float(max_range)}, inputs={"X": var, 'Scale': scale}, outputs={"Out": dequant_var}, ) return dequant_var