def LastValueQuantize(inputs, per_channel=False, init_min=-6.0, init_max=6.0, vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, name_prefix='LastValueQuant', reuse=None, is_training=True, num_bits=8, narrow_range=False): """Adds a layer that collects quantization ranges as last input ranges. LastValueQuantize creates variables called 'min' and 'max', representing the interval used for quantization and clamping. Args: inputs: a tensor containing values to be quantized. per_channel: (Optional) a boolean specifying whether to use different quantization ranges per output channel. init_min: a float scalar, the initial value for variable min. init_max: a float scalar, the initial value for variable max. vars_collection: (Optional) collection where to store variables for quantization interval ends. name_prefix: name_prefix for created nodes. reuse: whether or not the layer and its variables should be reused. To be able to reuse the layer scope must be given. is_training: Whether the op is applied to a training or eval graph. num_bits: Number of bits to use for quantization, must be between 2 and 8. narrow_range: Whether to use the narrow quantization range [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1]. Returns: a tensor containing quantized values. """ with variable_scope.variable_scope( None, default_name=name_prefix, values=[inputs], reuse=reuse): input_shape = inputs.get_shape() input_dim = len(input_shape) if per_channel: # Only support quantizing 1-, 2- and 4-dimensional tensors. assert input_dim in [1, 2, 4], ('Expected 1D, 2D or 4D input, was: %s in ' ' scope: %s' % (input_shape, name_prefix)) min_max_shape = [input_shape[-1]] else: min_max_shape = [] min_var = model_variable( 'min', shape=min_max_shape, initializer=init_ops.constant_initializer(init_min), collections=[vars_collection], trainable=False) max_var = model_variable( 'max', shape=min_max_shape, initializer=init_ops.constant_initializer(init_max), collections=[vars_collection], trainable=False) if not is_training: return _FakeQuantWithMinMaxVars( inputs, min_var, max_var, per_channel=per_channel, num_bits=num_bits, narrow_range=narrow_range) if per_channel: if input_dim == 2: reduce_dims = [0] elif input_dim == 4: reduce_dims = [0, 1, 2] if per_channel: if input_dim >= 2: batch_min = math_ops.reduce_min( inputs, reduction_indices=reduce_dims, name='BatchMin') else: batch_min = inputs else: batch_min = math_ops.reduce_min(inputs, name='BatchMin') # TFLite requires that 0.0 if always in the [min; max] range. batch_min = math_ops.minimum(batch_min, 0.0) assign_min = state_ops.assign(min_var, batch_min, name='AssignMinLast') if per_channel: if input_dim >= 2: batch_max = math_ops.reduce_max( inputs, reduction_indices=reduce_dims, name='BatchMax') else: batch_max = inputs else: batch_max = math_ops.reduce_max(inputs, name='BatchMax') # TFLite requires that 0.0 if always in the [min; max] range. batch_max = math_ops.maximum(batch_max, 0.0) assign_max = state_ops.assign(max_var, batch_max, name='AssignMaxLast') return _FakeQuantWithMinMaxVars( inputs, assign_min, assign_max, per_channel=per_channel, num_bits=num_bits, narrow_range=narrow_range)
def MovingAvgQuantize(inputs, per_channel=False, init_min=-6.0, init_max=6.0, ema_decay=0.999, vars_collection=ops.GraphKeys.MOVING_AVERAGE_VARIABLES, name_prefix='MovingAvgQuantize', reuse=None, is_training=True, num_bits=8, narrow_range=False): """Adds a layer that collects quantization ranges as EMAs of input ranges. MovingAvgQuantize creates variables called 'min' and 'max', representing the interval used for quantization and clamping. Args: inputs: a tensor containing values to be quantized. per_channel: (default False) a boolean specifying whether to use different quantization ranges per output channel. init_min: a float scalar, the initial value for variable min. init_max: a float scalar, the initial value for variable max. ema_decay: EMA decay parameter. vars_collection: (Optional) collection where to store variables for quantization interval ends. name_prefix: name_prefix for created nodes. reuse: whether or not the layer and its variables should be reused. To be able to reuse the layer scope must be given. is_training: Whether the op is applied to a training or eval graph. num_bits: Number of bits to use for quantization, must be between 2 and 8. narrow_range: Whether to use the narrow quantization range [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1]. Returns: a tensor containing quantized values. """ with variable_scope.variable_scope(None, default_name=name_prefix, values=[inputs], reuse=reuse) as scope: scope.set_partitioner(None) input_shape = inputs.get_shape() input_dim = len(input_shape) if per_channel: # Only support quantizing 1-, 2- and 4-dimensional tensors. assert input_dim in [1, 2, 4 ], ('Expected 1D, 2D or 4D input, was: %s in ' ' scope: %s' % (input_shape, name_prefix)) min_max_shape = [input_shape[-1]] else: min_max_shape = [] min_var = model_variable( 'min', shape=min_max_shape, initializer=init_ops.constant_initializer(init_min), collections=[vars_collection], trainable=False) max_var = model_variable( 'max', shape=min_max_shape, initializer=init_ops.constant_initializer(init_max), collections=[vars_collection], trainable=False) # modified by sf min_po2_var = model_variable( 'min_po2', shape=min_max_shape, initializer=init_ops.constant_initializer(init_min), collections=[vars_collection], trainable=False) max_po2_var = model_variable( 'max_po2', shape=min_max_shape, initializer=init_ops.constant_initializer(init_max), collections=[vars_collection], trainable=False) if not is_training: return _FakeQuantWithMinMaxVars( inputs, #min_var, #max_var, min_po2_var, max_po2_var, per_channel=per_channel, num_bits=num_bits, narrow_range=narrow_range) # end modified by sf if per_channel: if input_dim == 2: reduce_dims = [0] elif input_dim == 4: reduce_dims = [0, 1, 2] if per_channel: if input_dim >= 2: batch_min = math_ops.reduce_min(inputs, reduction_indices=reduce_dims, name='BatchMin') else: batch_min = inputs else: batch_min = math_ops.reduce_min(inputs, name='BatchMin') # B-eng requires that 0.0 if always in the [min; max] range. # modified by sf # ensure batch_min is not 0 batch_min = math_ops.minimum(batch_min, -0.0000001) # end modified by sf assign_min = moving_averages.assign_moving_average(min_var, batch_min, ema_decay, name='AssignMinEma') if per_channel: if input_dim >= 2: batch_max = math_ops.reduce_max(inputs, reduction_indices=reduce_dims, name='BatchMax') else: batch_max = inputs else: batch_max = math_ops.reduce_max(inputs, name='BatchMax') # B-eng requires that 0.0 if always in the [min; max] range. # modified by sf # ensure batch_max is not 0 batch_max = math_ops.maximum(batch_max, 0.0000001) # end modified by sf assign_max = moving_averages.assign_moving_average(max_var, batch_max, ema_decay, name='AssignMaxEma') # modified by sf min_log = tf.ceil(tf.log(tf.abs(assign_min)) / tf.log(2.0)) max_log = tf.ceil(tf.log(tf.abs(assign_max)) / tf.log(2.0)) minmax_log = math_ops.maximum(min_log, max_log) min_po2 = tf.negative(tf.pow(2.0, minmax_log)) max_po2 = tf.pow(2.0, minmax_log) - tf.pow(2.0, minmax_log - 7) assign_min_po2 = state_ops.assign(min_po2_var, min_po2, name='AssignMinPo2MovAvg') assign_max_po2 = state_ops.assign(max_po2_var, max_po2, name='AssignMaxPo2MovAvg') return _FakeQuantWithMinMaxVars( inputs, #assign_min, #assign_max, assign_min_po2, assign_max_po2, # end modified by sf per_channel=per_channel, num_bits=num_bits, narrow_range=narrow_range)