Internal change

PiperOrigin-RevId: 439776069

Author: daverim

tensorflow_model_optimization/python/core/quantization/keras/quantize.py

@@ -380,28 +380,33 @@ def _unwrap(layer):
 
     unwrapped_model = keras.models.clone_model(
         model_to_unwrap, input_tensors=None, clone_function=_unwrap)
-
     return unwrapped_model, layer_quantize_map, requires_output_quantize
 
   def _quantize(layer):  # pylint: disable=missing-docstring
-    if ((layer.name not in layer_quantize_map and
-         layer.name not in requires_output_quantize) or
-        (isinstance(layer, quantize_wrapper.QuantizeWrapper))):
-      # It supports for custom QuantizeWrapper.
-      return layer
-
+    # Handle quantize layer before any layers.
     # layer is a QuantizeLayer, possibly rebuild
     # layer with modified config from parameters stored in the map.
     if isinstance(layer, quantize_layer.QuantizeLayer):
       # If there is more than one usage of the input, even if all are concat,
       # we need to quantize.
       if len(layer._outbound_nodes) > 1:  # pylint: disable=protected-access
         return layer
+
       layer_config = layer.get_config()
+      if layer.name not in layer_quantize_map:  # Possibly added manually.
+        with quantize_scope():
+          return quantize_layer.QuantizeLayer.from_config(layer_config)
+
       for key, value in layer_quantize_map[layer.name].items():
         layer_config[key] = value
       return quantize_layer.QuantizeLayer.from_config(layer_config)
 
+    if ((layer.name not in layer_quantize_map and
+         layer.name not in requires_output_quantize) or
+        (isinstance(layer, quantize_wrapper.QuantizeWrapper))):
+      # It supports for custom QuantizeWrapper.
+      return layer
+
     if layer.name in requires_output_quantize:
       if not quantize_registry.supports(layer):
         return layer

tensorflow_model_optimization/python/core/quantization/keras/quantizers.py

@@ -129,16 +129,22 @@ def from_config(cls, config):
 class _QuantizeHelper(object):
   """Mixin with helper functions for quantizers."""
 
-  def _add_range_weights(self, layer, name):
+  def _add_range_weights(self, layer, name, per_axis=False, tensor_shape=None):
     """Add min and max vars to layer."""
+    shape = None
+    if per_axis and tensor_shape is not None:
+      shape = (tensor_shape[-1])
+
     min_weight = layer.add_weight(
         name + '_min',
         initializer=keras.initializers.Constant(-6.0),
-        trainable=False)
+        trainable=False,
+        shape=shape)
     max_weight = layer.add_weight(
         name + '_max',
         initializer=keras.initializers.Constant(6.0),
-        trainable=False)
+        trainable=False,
+        shape=shape)
 
     return {'min_var': min_weight, 'max_var': max_weight}
 
@@ -169,7 +175,7 @@ def __init__(self, num_bits, per_axis, symmetric, narrow_range):
     self.narrow_range = narrow_range
 
   def build(self, tensor_shape, name, layer):
-    return self._add_range_weights(layer, name)
+    return self._add_range_weights(layer, name, self.per_axis, tensor_shape)
 
   def __call__(self, inputs, training, weights, **kwargs):
     """Quantize tensor.