Create Quantizers specific to Conv/DConv for new quant scheme.

These new quantizers construct min/max vectors based on the shape
of the kernels in Conv/DConv, and support multiple scales per-channel
as in the new quant scheme.

They still need to be plugged into the layers, and also DConv special
case handling is not introduced yet.

PiperOrigin-RevId: 278964140

Author: nutsiepully

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

@@ -99,21 +99,21 @@ def LastValueQuantize(inputs,
 
     if per_channel:
       if input_dim >= 2:
-        batch_min = math_ops.reduce_min(
+        batch_min = math_ops.reduce_min_v1(
             inputs, reduction_indices=reduce_dims, name='BatchMin')
       else:
         batch_min = inputs
     else:
-      batch_min = math_ops.reduce_min(inputs, name='BatchMin')
+      batch_min = math_ops.reduce_min_v1(inputs, name='BatchMin')
 
     if per_channel:
       if input_dim >= 2:
-        batch_max = math_ops.reduce_max(
+        batch_max = math_ops.reduce_max_v1(
             inputs, reduction_indices=reduce_dims, name='BatchMax')
       else:
         batch_max = inputs
     else:
-      batch_max = math_ops.reduce_max(inputs, name='BatchMax')
+      batch_max = math_ops.reduce_max_v1(inputs, name='BatchMax')
 
     if symmetric:
       if narrow_range:

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

@@ -100,7 +100,8 @@ def __init__(self, num_bits, per_axis, symmetric, narrow_range):
 
     Args:
       num_bits: Number of bits for quantization
-      per_axis: Whether to apply per_axis quantization.
+      per_axis: Whether to apply per_axis quantization. The last dimension is
+        used as the axis.
       symmetric: If true, use symmetric quantization limits instead of training
         the minimum and maximum of each quantization range separately.
       narrow_range: In case of 8 bits, narrow_range nudges the quantized range
@@ -167,7 +168,8 @@ def __init__(self, num_bits, per_axis, symmetric, narrow_range):
 
     Args:
       num_bits: Number of bits for quantization
-      per_axis: Whether to apply per_axis quantization.
+      per_axis: Whether to apply per_axis quantization. The last dimension is
+        used as the axis.
       symmetric: If true, use symmetric quantization limits instead of training
         the minimum and maximum of each quantization range separately.
       narrow_range: In case of 8 bits, narrow_range nudges the quantized range

tensorflow_model_optimization/python/core/quantization/keras/tflite/BUILD

@@ -4,6 +4,36 @@ package(default_visibility = [
 
 licenses(["notice"])  # Apache 2.0
 
+py_library(
+    name = "tflite_quantizers",
+    srcs = [
+        "tflite_quantizers.py",
+    ],
+    srcs_version = "PY2AND3",
+    visibility = ["//visibility:public"],
+    deps = [
+        # tensorflow dep1,
+        # python/keras tensorflow dep2,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+    ],
+)
+
+py_test(
+    name = "tflite_quantizers_test",
+    srcs = [
+        "tflite_quantizers_test.py",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":tflite_quantizers",
+        # absl/testing:parameterized dep1,
+        # tensorflow dep1,
+        # python/keras tensorflow dep2,
+    ],
+)
+
 py_library(
     name = "tflite_quantize_registry",
     srcs = [

tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_quantizers.py

@@ -0,0 +1,49 @@
+# Copyright 2019 The TensorFlow 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.
+# ==============================================================================
+"""Quantizers specific to TFLite.
+
+Module: tfmot.quantization.keras.tflite
+"""
+
+from tensorflow.python.keras import initializers
+
+from tensorflow_model_optimization.python.core.quantization.keras import quantizers
+
+
+class ConvWeightsQuantizer(quantizers.LastValueQuantizer):
+  """Quantizer for handling weights in Conv2D/DepthwiseConv2D layers."""
+
+  def __init__(self):
+    """Construct LastValueQuantizer with params specific for TFLite Convs."""
+
+    super(ConvWeightsQuantizer, self).__init__(
+        num_bits=8,
+        per_axis=True,
+        symmetric=True,
+        narrow_range=True)
+
+  def build(self, tensor_shape, name, layer):
+    min_weight = layer.add_weight(
+        name + '_min',
+        shape=(tensor_shape[-1],),
+        initializer=initializers.Constant(-6.0),
+        trainable=False,)
+    max_weight = layer.add_weight(
+        name + '_max',
+        shape=(tensor_shape[-1],),
+        initializer=initializers.Constant(6.0),
+        trainable=False)
+
+    return [min_weight, max_weight]

tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_quantizers_test.py

@@ -0,0 +1,62 @@
+# Copyright 2019 The TensorFlow 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.
+# ==============================================================================
+"""Tests for TFLite Quantizers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+from tensorflow.python import keras
+from tensorflow.python.platform import test
+
+from tensorflow_model_optimization.python.core.quantization.keras.tflite import tflite_quantizers
+
+ConvWeightsQuantizer = tflite_quantizers.ConvWeightsQuantizer
+
+
+class ConvWeightsQuantizerTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      (keras.layers.Conv2D, {
+          'filters': 5,
+          'kernel_size': (2, 2)
+      }),
+      (keras.layers.DepthwiseConv2D, {
+          'kernel_size': (2, 2),
+          'depth_multiplier': 5,
+      })
+  )
+  def testConstructsMinMaxVarsCorrectShape(self, layer_type, kwargs):
+    quantizer = ConvWeightsQuantizer()
+
+    model = keras.Sequential([
+        layer_type(input_shape=(5, 2, 3), **kwargs)])
+    layer = model.layers[0]
+
+    min_var, max_var = quantizer.build(
+        layer.weights[0].shape, 'kernel', layer)
+    # TODO(pulkitb): Add value test to ensure per-axis quantization is
+    # happening properly. Probably to quant_ops_test.py
+    quantized_weight = quantizer(layer.weights[0], 0, True,  # pylint: disable=unused-variable
+                                 **{'min_var': min_var, 'max_var': max_var})
+
+    self.assertEqual(5, min_var.shape)
+    self.assertEqual(5, max_var.shape)
+
+
+if __name__ == '__main__':
+  test.main()