def _quantize(layer): # pylint: disable=missing-docstring if (layer.name not in layer_quantize_map and layer.name not in requires_output_quantize): return layer if layer.name in requires_output_quantize: if not quantize_registry.supports(layer): return layer full_quantize_config = quantize_registry.get_quantize_config(layer) if not full_quantize_config: return layer quantize_config = OutputOnlyConfig(full_quantize_config) else: quantize_config = layer_quantize_map[layer.name].get('quantize_config') if not quantize_config and quantize_registry.supports(layer): quantize_config = quantize_registry.get_quantize_config(layer) if not quantize_config: error_msg = ( 'Layer {}:{} is not supported. You can quantize this ' 'layer by passing a `tfmot.quantization.keras.QuantizeConfig` ' 'instance to the `quantize_annotate_layer` ' 'API.') raise RuntimeError( error_msg.format(layer.name, layer.__class__, quantize_registry.__class__)) # `QuantizeWrapper` does not copy any additional layer params from # `QuantizeAnnotate`. This should generally be fine, but occasionally # `QuantizeAnnotate` wrapper may contain `batch_input_shape` like params. # TODO(pulkitb): Ensure this does not affect model cloning. return quantize_wrapper.QuantizeWrapper(layer, quantize_config)
def quantize_layer(layer, apply_quantization=True, quantize_config=None): """Quantizes a layer. It is useful for quantization aware training Args: layer: input layer to quantize apply_quantization: if True layer is quantized, otherwise not quantize_config: quantization config for special cases such as sequence of convolution and batch normalization Returns: quantized layer """ if apply_quantization: scheme = tfmot.quantization.keras.default_8bit.Default8BitQuantizeScheme( ) quantize_registry = scheme.get_quantize_registry() if not quantize_registry.supports(layer): logging.info('layer is not supported: %s', str(layer)) return layer if quantize_config is None: quantize_config = quantize_registry.get_quantize_config(layer) return quantize_wrapper.QuantizeWrapper(layer, quantize_config) else: return layer
def _quantize(layer): # pylint: disable=missing-docstring if layer.name not in layer_quantize_map: return layer quantize_config = layer_quantize_map[layer.name].get('quantize_config') if not quantize_config and quantize_registry.supports(layer): quantize_config = quantize_registry.get_quantize_config(layer) if not quantize_config: error_msg = ('Could not find a suitable QuantizeConfig for layer {}. ' 'Either the registry {} should be provide one, or the user ' 'should provide one while annotating the layer using ' 'QuantizeAnnotate.') raise RuntimeError(error_msg.format( layer.__class__, quantize_registry.__class__)) # `QuantizeWrapper` does not copy any additional layer params from # `QuantizeAnnotate`. This should generally be fine, but occasionally # `QuantizeAnnotate` wrapper may contain `batch_input_shape` like params. # TODO(pulkitb): Ensure this does not affect model cloning. return quantize_wrapper.QuantizeWrapper(layer, quantize_config)
def quantize_layer(layer, apply_quantization=True): """Quantizes a layer. It is useful for quantization aware training Args: layer: input layer to quantize apply_quantization: if True layer is quantized, otherwise not returned Returns: quantized layer """ if apply_quantization: scheme = tfmot.quantization.keras.default_8bit.Default8BitQuantizeScheme( ) quantize_registry = scheme.get_quantize_registry() if not quantize_registry.supports(layer): return layer quantize_config = quantize_registry.get_quantize_config(layer) return quantize_wrapper.QuantizeWrapper(layer, quantize_config) else: return layer
def quantize(layer, quantize_config=None): if quantize_config is None: quantize_config = default_8bit_quantize_configs.Default8BitOutputQuantizeConfig( ) return quantize_wrapper.QuantizeWrapper(layer, quantize_config=quantize_config)