Merge branch 'tensorflow:master' into panoptic-deeplab-modeling

official/projects/qat/vision/modeling/layers/nn_blocks_test.py

+# Copyright 2022 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.
+
+# Lint as: python3
+"""Tests for nn_blocks."""
+
+from typing import Any, Iterable, Tuple
+# Import libraries
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from official.projects.qat.vision.modeling.layers import nn_blocks
+
+
+def distribution_strategy_combinations() -> Iterable[Tuple[Any, ...]]:
+  """Returns the combinations of end-to-end tests to run."""
+  return combinations.combine(
+      distribution=[
+          strategy_combinations.default_strategy,
+          strategy_combinations.cloud_tpu_strategy,
+          strategy_combinations.one_device_strategy_gpu,
+      ],
+  )
+
+
+class NNBlocksTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (nn_blocks.BottleneckBlockQuantized, 1, False, 0.0, None),
+      (nn_blocks.BottleneckBlockQuantized, 2, True, 0.2, 0.25),
+  )
+  def test_bottleneck_block_creation(self, block_fn, strides, use_projection,
+                                     stochastic_depth_drop_rate, se_ratio):
+    input_size = 128
+    filter_size = 256
+    inputs = tf.keras.Input(
+        shape=(input_size, input_size, filter_size * 4), batch_size=1)
+    block = block_fn(
+        filter_size,
+        strides,
+        use_projection=use_projection,
+        se_ratio=se_ratio,
+        stochastic_depth_drop_rate=stochastic_depth_drop_rate)
+
+    features = block(inputs)
+
+    self.assertAllEqual(
+        [1, input_size // strides, input_size // strides, filter_size * 4],
+        features.shape.as_list())
+
+  @parameterized.parameters(
+      (nn_blocks.InvertedBottleneckBlockQuantized, 1, 1, None, None),
+      (nn_blocks.InvertedBottleneckBlockQuantized, 6, 1, None, None),
+      (nn_blocks.InvertedBottleneckBlockQuantized, 1, 2, None, None),
+      (nn_blocks.InvertedBottleneckBlockQuantized, 1, 1, 0.2, None),
+      (nn_blocks.InvertedBottleneckBlockQuantized, 1, 1, None, 0.2),
+  )
+  def test_invertedbottleneck_block_creation(
+      self, block_fn, expand_ratio, strides, se_ratio,
+      stochastic_depth_drop_rate):
+    input_size = 128
+    in_filters = 24
+    out_filters = 40
+    inputs = tf.keras.Input(
+        shape=(input_size, input_size, in_filters), batch_size=1)
+    block = block_fn(
+        in_filters=in_filters,
+        out_filters=out_filters,
+        expand_ratio=expand_ratio,
+        strides=strides,
+        se_ratio=se_ratio,
+        stochastic_depth_drop_rate=stochastic_depth_drop_rate,
+        output_intermediate_endpoints=False)
+
+    features = block(inputs)
+
+    self.assertAllEqual(
+        [1, input_size // strides, input_size // strides, out_filters],
+        features.shape.as_list())
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/qat/vision/modeling/layers/nn_layers_test.py

+# Copyright 2022 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 nn_layers."""
+
+# Import libraries
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.projects.qat.vision.modeling.layers import nn_layers
+
+
+class NNLayersTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      ('deeplabv3plus', 1),
+      ('deeplabv3plus', 2),
+      ('deeplabv3', 1),
+      ('deeplabv3', 2),
+  )
+  def test_segmentation_head_creation(self, feature_fusion, upsample_factor):
+    input_size = 128
+    decoder_outupt_size = input_size // 2
+
+    decoder_output = tf.random.uniform(
+        (2, decoder_outupt_size, decoder_outupt_size, 64), dtype=tf.float32)
+    backbone_output = tf.random.uniform((2, input_size, input_size, 32),
+                                        dtype=tf.float32)
+    segmentation_head = nn_layers.SegmentationHeadQuantized(
+        num_classes=5,
+        level=4,
+        upsample_factor=upsample_factor,
+        low_level=2,
+        low_level_num_filters=128,
+        feature_fusion=feature_fusion)
+
+    features = segmentation_head((backbone_output, decoder_output))
+
+    expected_shape = (
+        input_size
+        if feature_fusion == 'deeplabv3plus' else decoder_outupt_size)
+    self.assertAllEqual([
+        2, expected_shape * upsample_factor, expected_shape * upsample_factor, 5
+    ], features.shape.as_list())
+
+  @parameterized.parameters(
+      (None, []),
+      (None, [6, 12, 18]),
+      ([32, 32], [6, 12, 18]),
+  )
+  def test_spatial_pyramid_pooling_creation(self, pool_kernel_size,
+                                            dilation_rates):
+    inputs = tf.keras.Input(shape=(64, 64, 128), dtype=tf.float32)
+    layer = nn_layers.SpatialPyramidPoolingQuantized(
+        output_channels=256,
+        dilation_rates=dilation_rates,
+        pool_kernel_size=pool_kernel_size)
+    output = layer(inputs)
+    self.assertAllEqual([None, 64, 64, 256], output.shape)
+
+  @parameterized.parameters(
+      (3, [6, 12, 18, 24], 128),
+      (3, [6, 12, 18], 128),
+      (3, [6, 12], 256),
+      (4, [], 128),
+      (4, [6, 12, 18], 128),
+      (4, [], 256),
+  )
+  def test_aspp_creation(self, level, dilation_rates, num_filters):
+    input_size = 128 // 2**level
+    tf.keras.backend.set_image_data_format('channels_last')
+    endpoints = tf.random.uniform(
+        shape=(2, input_size, input_size, 64), dtype=tf.float32)
+
+    network = nn_layers.ASPPQuantized(
+        level=level, dilation_rates=dilation_rates, num_filters=num_filters)
+
+    feats = network(endpoints)
+
+    self.assertAllEqual([2, input_size, input_size, num_filters],
+                        feats.shape.as_list())
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/qat/vision/modeling/segmentation_model.py

+# Copyright 2022 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.
+
+"""Build segmentation models."""
+from typing import Any, Mapping, Union
+
+# Import libraries
+import tensorflow as tf
+
+layers = tf.keras.layers
+
+
+@tf.keras.utils.register_keras_serializable(package='Vision')
+class SegmentationModelQuantized(tf.keras.Model):
+  """A Segmentation class model.
+
+  Input images are passed through backbone first. Decoder network is then
+  applied, and finally, segmentation head is applied on the output of the
+  decoder network. Layers such as ASPP should be part of decoder. Any feature
+  fusion is done as part of the segmentation head (i.e. deeplabv3+ feature
+  fusion is not part of the decoder, instead it is part of the segmentation
+  head). This way, different feature fusion techniques can be combined with
+  different backbones, and decoders.
+  """
+
+  def __init__(self, backbone: tf.keras.Model, decoder: tf.keras.layers.Layer,
+               head: tf.keras.layers.Layer,
+               input_specs: tf.keras.layers.InputSpec, **kwargs):
+    """Segmentation initialization function.
+
+    Args:
+      backbone: a backbone network.
+      decoder: a decoder network. E.g. FPN.
+      head: segmentation head.
+      input_specs: The shape specifications of input tensor.
+      **kwargs: keyword arguments to be passed.
+    """
+    inputs = tf.keras.Input(shape=input_specs.shape[1:], name=input_specs.name)
+    backbone_features = backbone(inputs)
+
+    if decoder:
+      backbone_feature = backbone_features[str(decoder.get_config()['level'])]
+      decoder_feature = decoder(backbone_feature)
+    else:
+      decoder_feature = backbone_features
+
+    backbone_feature = backbone_features[str(head.get_config()['low_level'])]
+    x = {'logits': head((backbone_feature, decoder_feature))}
+    super().__init__(inputs=inputs, outputs=x, **kwargs)
+    self._config_dict = {
+        'backbone': backbone,
+        'decoder': decoder,
+        'head': head,
+    }
+    self.backbone = backbone
+    self.decoder = decoder
+    self.head = head
+
+  @property
+  def checkpoint_items(
+      self) -> Mapping[str, Union[tf.keras.Model, tf.keras.layers.Layer]]:
+    """Returns a dictionary of items to be additionally checkpointed."""
+    items = dict(backbone=self.backbone, head=self.head)
+    if self.decoder is not None:
+      items.update(decoder=self.decoder)
+    return items
+
+  def get_config(self) -> Mapping[str, Any]:
+    return self._config_dict
+
+  @classmethod
+  def from_config(cls, config, custom_objects=None):
+    return cls(**config)

official/projects/qat/vision/n_bit/__init__.py

+# Copyright 2022 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.
+
+# Lint as: python3
+"""Configs package definition."""
+
+from official.projects.qat.vision.n_bit import configs
+from official.projects.qat.vision.n_bit import schemes
+from official.projects.qat.vision.n_bit.nn_blocks import BottleneckBlockNBitQuantized
+from official.projects.qat.vision.n_bit.nn_blocks import Conv2DBNBlockNBitQuantized
+from official.projects.qat.vision.n_bit.nn_blocks import InvertedBottleneckBlockNBitQuantized

official/projects/qat/vision/n_bit/configs.py

+# Copyright 2022 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.
+
+"""Default 8-bit QuantizeConfigs."""
+from typing import Sequence, Callable, Tuple, Any, Dict
+
+import tensorflow as tf
+import tensorflow_model_optimization as tfmot
+
+
+Quantizer = tfmot.quantization.keras.quantizers.Quantizer
+Layer = tf.keras.layers.Layer
+Activation = Callable[[tf.Tensor], tf.Tensor]
+WeightAndQuantizer = Tuple[tf.Variable, Quantizer]
+ActivationAndQuantizer = Tuple[Activation, Quantizer]
+
+
+class DefaultNBitOutputQuantizeConfig(
+    tfmot.quantization.keras.QuantizeConfig):
+  """QuantizeConfig which only quantizes the output from a layer."""
+
+  def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8):
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+  def get_weights_and_quantizers(
+      self, layer: Layer) -> Sequence[WeightAndQuantizer]:
+    return []
+
+  def get_activations_and_quantizers(
+      self, layer: Layer) -> Sequence[ActivationAndQuantizer]:
+    return []
+
+  def set_quantize_weights(self,
+                           layer: Layer,
+                           quantize_weights: Sequence[tf.Tensor]):
+    pass
+
+  def set_quantize_activations(self,
+                               layer: Layer,
+                               quantize_activations: Sequence[Activation]):
+    pass
+
+  def get_output_quantizers(self, layer: Layer) -> Sequence[Quantizer]:
+    return [
+        tfmot.quantization.keras.quantizers.MovingAverageQuantizer(
+            num_bits=self._num_bits_activation, per_axis=False,
+            symmetric=False, narrow_range=False)  # activation/output
+    ]
+
+  def get_config(self) -> Dict[str, Any]:
+    return {
+        'num_bits_weight': self._num_bits_weight,
+        'num_bits_activation': self._num_bits_activation,
+    }
+
+
+class NoOpQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
+  """QuantizeConfig which does not quantize any part of the layer."""
+
+  def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8):
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+  def get_weights_and_quantizers(
+      self, layer: Layer) -> Sequence[WeightAndQuantizer]:
+    return []
+
+  def get_activations_and_quantizers(
+      self, layer: Layer) -> Sequence[ActivationAndQuantizer]:
+    return []
+
+  def set_quantize_weights(
+      self,
+      layer: Layer,
+      quantize_weights: Sequence[tf.Tensor]):
+    pass
+
+  def set_quantize_activations(
+      self,
+      layer: Layer,
+      quantize_activations: Sequence[Activation]):
+    pass
+
+  def get_output_quantizers(self, layer: Layer) -> Sequence[Quantizer]:
+    return []
+
+  def get_config(self) -> Dict[str, Any]:
+    return {
+        'num_bits_weight': self._num_bits_weight,
+        'num_bits_activation': self._num_bits_activation,
+    }
+
+
+class DefaultNBitQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
+  """QuantizeConfig for non recurrent Keras layers."""
+
+  def __init__(self,
+               weight_attrs: Sequence[str],
+               activation_attrs: Sequence[str],
+               quantize_output: bool,
+               num_bits_weight: int = 8,
+               num_bits_activation: int = 8):
+    """Initializes a default N-bit quantize config."""
+    self.weight_attrs = weight_attrs
+    self.activation_attrs = activation_attrs
+    self.quantize_output = quantize_output
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+    # TODO(pulkitb): For some layers such as Conv2D, per_axis should be True.
+    # Add mapping for which layers support per_axis.
+    self.weight_quantizer = tfmot.quantization.keras.quantizers.LastValueQuantizer(
+        num_bits=num_bits_weight, per_axis=False,
+        symmetric=True, narrow_range=True)  # weight
+    self.activation_quantizer = tfmot.quantization.keras.quantizers.MovingAverageQuantizer(
+        num_bits=num_bits_activation, per_axis=False,
+        symmetric=False, narrow_range=False)  # activation/output
+
+  def get_weights_and_quantizers(
+      self, layer: Layer) -> Sequence[WeightAndQuantizer]:
+    """See base class."""
+    return [(getattr(layer, weight_attr), self.weight_quantizer)
+            for weight_attr in self.weight_attrs]
+
+  def get_activations_and_quantizers(
+      self, layer: Layer) -> Sequence[ActivationAndQuantizer]:
+    """See base class."""
+    return [(getattr(layer, activation_attr), self.activation_quantizer)
+            for activation_attr in self.activation_attrs]
+
+  def set_quantize_weights(
+      self,
+      layer: Layer,
+      quantize_weights: Sequence[tf.Tensor]):
+    """See base class."""
+    if len(self.weight_attrs) != len(quantize_weights):
+      raise ValueError(
+          '`set_quantize_weights` called on layer {} with {} '
+          'weight parameters, but layer expects {} values.'.format(
+              layer.name, len(quantize_weights), len(self.weight_attrs)))
+
+    for weight_attr, weight in zip(self.weight_attrs, quantize_weights):
+      current_weight = getattr(layer, weight_attr)
+      if current_weight.shape != weight.shape:
+        raise ValueError('Existing layer weight shape {} is incompatible with'
+                         'provided weight shape {}'.format(
+                             current_weight.shape, weight.shape))
+
+      setattr(layer, weight_attr, weight)
+
+  def set_quantize_activations(
+      self,
+      layer: Layer,
+      quantize_activations: Sequence[Activation]):
+    """See base class."""
+    if len(self.activation_attrs) != len(quantize_activations):
+      raise ValueError(
+          '`set_quantize_activations` called on layer {} with {} '
+          'activation parameters, but layer expects {} values.'.format(
+              layer.name, len(quantize_activations),
+              len(self.activation_attrs)))
+
+    for activation_attr, activation in zip(
+        self.activation_attrs, quantize_activations):
+      setattr(layer, activation_attr, activation)
+
+  def get_output_quantizers(self, layer: Layer) -> Sequence[Quantizer]:
+    """See base class."""
+    if self.quantize_output:
+      return [self.activation_quantizer]
+    return []
+
+  @classmethod
+  def from_config(cls, config: Dict[str, Any]) -> object:
+    """Instantiates a `DefaultNBitQuantizeConfig` from its config.
+
+    Args:
+        config: Output of `get_config()`.
+
+    Returns:
+        A `DefaultNBitQuantizeConfig` instance.
+    """
+    return cls(**config)
+
+  def get_config(self) -> Dict[str, Any]:
+    """Get a config for this quantize config."""
+    # TODO(pulkitb): Add weight and activation quantizer to config.
+    # Currently it's created internally, but ideally the quantizers should be
+    # part of the constructor and passed in from the registry.
+    return {
+        'weight_attrs': self.weight_attrs,
+        'activation_attrs': self.activation_attrs,
+        'quantize_output': self.quantize_output,
+        'num_bits_weight': self._num_bits_weight,
+        'num_bits_activation': self._num_bits_activation
+    }
+
+  def __eq__(self, other):
+    if not isinstance(other, DefaultNBitQuantizeConfig):
+      return False
+
+    return (self.weight_attrs == other.weight_attrs and
+            self.activation_attrs == self.activation_attrs and
+            self.weight_quantizer == other.weight_quantizer and
+            self.activation_quantizer == other.activation_quantizer and
+            self.quantize_output == other.quantize_output)
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
+
+class DefaultNBitConvWeightsQuantizer(
+    tfmot.quantization.keras.quantizers.LastValueQuantizer):
+  """Quantizer for handling weights in Conv2D/DepthwiseConv2D layers."""
+
+  def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8):
+    """Construct LastValueQuantizer with params specific for TFLite Convs."""
+
+    super(DefaultNBitConvWeightsQuantizer, self).__init__(
+        num_bits=num_bits_weight, per_axis=True,
+        symmetric=True, narrow_range=True)  # weight
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+  def build(self,
+            tensor_shape: tf.TensorShape,
+            name: str,
+            layer: Layer):
+    """Build min/max quantization variables."""
+    min_weight = layer.add_weight(
+        name + '_min',
+        shape=(tensor_shape[-1],),
+        initializer=tf.keras.initializers.Constant(-6.0),
+        trainable=False)
+    max_weight = layer.add_weight(
+        name + '_max',
+        shape=(tensor_shape[-1],),
+        initializer=tf.keras.initializers.Constant(6.0),
+        trainable=False)
+
+    return {'min_var': min_weight, 'max_var': max_weight}
+
+
+class NoQuantizer(tfmot.quantization.keras.quantizers.Quantizer):
+  """Dummy quantizer for explicitly not quantize."""
+
+  def __call__(self, inputs, training, weights, **kwargs):
+    return tf.identity(inputs)
+
+  def get_config(self):
+    return {}
+
+  def build(self, tensor_shape, name, layer):
+    return {}
+
+
+class DefaultNBitConvQuantizeConfig(DefaultNBitQuantizeConfig):
+  """QuantizeConfig for Conv2D/DepthwiseConv2D layers."""
+
+  def __init__(self,
+               weight_attrs: Sequence[str],
+               activation_attrs: Sequence[str],
+               quantize_output: bool,
+               num_bits_weight: int = 8,
+               num_bits_activation: int = 8):
+    """Initializes default N-bit quantization config for the conv layer."""
+    super().__init__(weight_attrs=weight_attrs,
+                     activation_attrs=activation_attrs,
+                     quantize_output=quantize_output,
+                     num_bits_weight=num_bits_weight,
+                     num_bits_activation=num_bits_activation)
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+    self.weight_quantizer = DefaultNBitConvWeightsQuantizer(
+        num_bits_weight=num_bits_weight,
+        num_bits_activation=num_bits_activation)
+
+
+class DefaultNBitActivationQuantizeConfig(
+    tfmot.quantization.keras.QuantizeConfig):
+  """QuantizeConfig for keras.layers.Activation.
+
+  `keras.layers.Activation` needs a separate `QuantizeConfig` since the
+  decision to quantize depends on the specific activation type.
+  """
+
+  def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8):
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+  def _assert_activation_layer(self, layer: Layer):
+    if not isinstance(layer, tf.keras.layers.Activation):
+      raise RuntimeError(
+          'DefaultNBitActivationQuantizeConfig can only be used with '
+          '`keras.layers.Activation`.')
+
+  def get_weights_and_quantizers(
+      self, layer: Layer) -> Sequence[WeightAndQuantizer]:
+    """See base class."""
+    self._assert_activation_layer(layer)
+    return []
+
+  def get_activations_and_quantizers(
+      self, layer: Layer) -> Sequence[ActivationAndQuantizer]:
+    """See base class."""
+    self._assert_activation_layer(layer)
+    return []
+
+  def set_quantize_weights(
+      self,
+      layer: Layer,
+      quantize_weights: Sequence[tf.Tensor]):
+    """See base class."""
+    self._assert_activation_layer(layer)
+
+  def set_quantize_activations(
+      self,
+      layer: Layer,
+      quantize_activations: Sequence[Activation]):
+    """See base class."""
+    self._assert_activation_layer(layer)
+
+  def get_output_quantizers(self, layer: Layer) -> Sequence[Quantizer]:
+    """See base class."""
+    self._assert_activation_layer(layer)
+
+    if not hasattr(layer.activation, '__name__'):
+      raise ValueError('Activation {} not supported by '
+                       'DefaultNBitActivationQuantizeConfig.'.format(
+                           layer.activation))
+
+    # This code is copied from TFMOT repo, but added relu6 to support mobilenet.
+    if layer.activation.__name__ in ['relu', 'relu6', 'swish']:
+      # 'relu' should generally get fused into the previous layer.
+      return [tfmot.quantization.keras.quantizers.MovingAverageQuantizer(
+          num_bits=self._num_bits_activation, per_axis=False,
+          symmetric=False, narrow_range=False)]  # activation/output
+    elif layer.activation.__name__ in ['linear', 'softmax', 'sigmoid']:
+      return []
+
+    raise ValueError('Activation {} not supported by '
+                     'DefaultNBitActivationQuantizeConfig.'.format(
+                         layer.activation))
+
+  def get_config(self) -> Dict[str, Any]:
+    """Get a config for this quantizer config."""
+    return {
+        'num_bits_weight': self._num_bits_weight,
+        'num_bits_activation': self._num_bits_activation,
+    }
+
+
+def _types_dict():
+  return {
+      'DefaultNBitOutputQuantizeConfig':
+          DefaultNBitOutputQuantizeConfig,
+      'NoOpQuantizeConfig':
+          NoOpQuantizeConfig,
+      'DefaultNBitQuantizeConfig':
+          DefaultNBitQuantizeConfig,
+      'DefaultNBitConvWeightsQuantizer':
+          DefaultNBitConvWeightsQuantizer,
+      'DefaultNBitConvQuantizeConfig':
+          DefaultNBitConvQuantizeConfig,
+      'DefaultNBitActivationQuantizeConfig':
+          DefaultNBitActivationQuantizeConfig,
+  }

official/projects/qat/vision/n_bit/configs_test.py

+# Copyright 2022 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 configs.py."""
+
+# Import libraries
+
+import numpy as np
+import tensorflow as tf
+
+import tensorflow_model_optimization as tfmot
+
+from official.projects.qat.vision.n_bit import configs
+
+
+class _TestHelper(object):
+
+  def _convert_list(self, list_of_tuples):
+    """Transforms a list of 2-tuples to a tuple of 2 lists.
+
+    `QuantizeConfig` methods return a list of 2-tuples in the form
+    [(weight1, quantizer1), (weight2, quantizer2)]. This function converts
+    it into a 2-tuple of lists. ([weight1, weight2]), (quantizer1, quantizer2).
+
+    Args:
+      list_of_tuples: List of 2-tuples.
+
+    Returns:
+      2-tuple of lists.
+    """
+    list1 = []
+    list2 = []
+    for a, b in list_of_tuples:
+      list1.append(a)
+      list2.append(b)
+
+    return list1, list2
+
+  # TODO(pulkitb): Consider asserting on full equality for quantizers.
+
+  def _assert_weight_quantizers(self, quantizer_list):
+    for quantizer in quantizer_list:
+      self.assertIsInstance(
+          quantizer,
+          tfmot.quantization.keras.quantizers.LastValueQuantizer)
+
+  def _assert_activation_quantizers(self, quantizer_list):
+    for quantizer in quantizer_list:
+      self.assertIsInstance(
+          quantizer,
+          tfmot.quantization.keras.quantizers.MovingAverageQuantizer)
+
+  def _assert_kernel_equality(self, a, b):
+    self.assertAllEqual(a.numpy(), b.numpy())
+
+
+class DefaultNBitQuantizeConfigTest(tf.test.TestCase, _TestHelper):
+
+  def _simple_dense_layer(self):
+    layer = tf.keras.layers.Dense(2)
+    layer.build(input_shape=(3,))
+    return layer
+
+  def testGetsQuantizeWeightsAndQuantizers(self):
+    layer = self._simple_dense_layer()
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+    (weights, weight_quantizers) = self._convert_list(
+        quantize_config.get_weights_and_quantizers(layer))
+
+    self._assert_weight_quantizers(weight_quantizers)
+    self.assertEqual([layer.kernel], weights)
+
+  def testGetsQuantizeActivationsAndQuantizers(self):
+    layer = self._simple_dense_layer()
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+    (activations, activation_quantizers) = self._convert_list(
+        quantize_config.get_activations_and_quantizers(layer))
+
+    self._assert_activation_quantizers(activation_quantizers)
+    self.assertEqual([layer.activation], activations)
+
+  def testSetsQuantizeWeights(self):
+    layer = self._simple_dense_layer()
+    quantize_kernel = tf.keras.backend.variable(
+        np.ones(layer.kernel.shape.as_list()))
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+    quantize_config.set_quantize_weights(layer, [quantize_kernel])
+
+    self._assert_kernel_equality(layer.kernel, quantize_kernel)
+
+  def testSetsQuantizeActivations(self):
+    layer = self._simple_dense_layer()
+    quantize_activation = tf.keras.activations.relu
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+    quantize_config.set_quantize_activations(layer, [quantize_activation])
+
+    self.assertEqual(layer.activation, quantize_activation)
+
+  def testSetsQuantizeWeights_ErrorOnWrongNumberOfWeights(self):
+    layer = self._simple_dense_layer()
+    quantize_kernel = tf.keras.backend.variable(
+        np.ones(layer.kernel.shape.as_list()))
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+
+    with self.assertRaises(ValueError):
+      quantize_config.set_quantize_weights(layer, [])
+
+    with self.assertRaises(ValueError):
+      quantize_config.set_quantize_weights(layer,
+                                           [quantize_kernel, quantize_kernel])
+
+  def testSetsQuantizeWeights_ErrorOnWrongShapeOfWeight(self):
+    layer = self._simple_dense_layer()
+    quantize_kernel = tf.keras.backend.variable(np.ones([1, 2]))
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+
+    with self.assertRaises(ValueError):
+      quantize_config.set_quantize_weights(layer, [quantize_kernel])
+
+  def testSetsQuantizeActivations_ErrorOnWrongNumberOfActivations(self):
+    layer = self._simple_dense_layer()
+    quantize_activation = tf.keras.activations.relu
+    num_bits_weight = 4
+    num_bits_activation = 4
+
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+
+    with self.assertRaises(ValueError):
+      quantize_config.set_quantize_activations(layer, [])
+
+    with self.assertRaises(ValueError):
+      quantize_config.set_quantize_activations(
+          layer, [quantize_activation, quantize_activation])
+
+  def testGetsResultQuantizers_ReturnsQuantizer(self):
+    layer = self._simple_dense_layer()
+    num_bits_weight = 4
+    num_bits_activation = 4
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        [], [], True, num_bits_weight, num_bits_activation)
+
+    output_quantizers = quantize_config.get_output_quantizers(layer)
+
+    self.assertLen(output_quantizers, 1)
+    self._assert_activation_quantizers(output_quantizers)
+
+  def testGetsResultQuantizers_EmptyWhenFalse(self):
+    layer = self._simple_dense_layer()
+    num_bits_weight = 4
+    num_bits_activation = 4
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        [], [], False, num_bits_weight, num_bits_activation)
+
+    output_quantizers = quantize_config.get_output_quantizers(layer)
+
+    self.assertEqual([], output_quantizers)
+
+  def testSerialization(self):
+    num_bits_weight = 4
+    num_bits_activation = 4
+    quantize_config = configs.DefaultNBitQuantizeConfig(
+        ['kernel'], ['activation'], False, num_bits_weight, num_bits_activation)
+
+    expected_config = {
+        'class_name': 'DefaultNBitQuantizeConfig',
+        'config': {
+            'weight_attrs': ['kernel'],
+            'activation_attrs': ['activation'],
+            'quantize_output': False,
+            'num_bits_weight': 4,
+            'num_bits_activation': 4
+        }
+    }
+    serialized_quantize_config = tf.keras.utils.serialize_keras_object(
+        quantize_config)
+
+    self.assertEqual(expected_config, serialized_quantize_config)
+
+    quantize_config_from_config = tf.keras.utils.deserialize_keras_object(
+        serialized_quantize_config,
+        module_objects=globals(),
+        custom_objects=configs._types_dict())
+
+    self.assertEqual(quantize_config, quantize_config_from_config)
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/qat/vision/quantization/__init__.py

+# Copyright 2022 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.
+
+# Lint as: python3
+"""Configs package definition."""
+
+from official.projects.qat.vision.quantization import configs
+from official.projects.qat.vision.quantization import schemes

official/projects/qat/vision/registry_imports.py

+# Copyright 2022 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.
+
+"""All necessary imports for registration on qat project."""
+# pylint: disable=unused-import
+from official.projects.qat.vision import configs
+from official.projects.qat.vision.modeling import layers
+from official.projects.qat.vision.tasks import image_classification
+from official.projects.qat.vision.tasks import retinanet
+from official.projects.qat.vision.tasks import semantic_segmentation