Internal change

PiperOrigin-RevId: 404080616

official/projects/edgetpu/vision/modeling/__init__.py

+# Copyright 2021 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.
+

official/projects/edgetpu/vision/modeling/backbones/__init__.py

+# Copyright 2021 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.
+

official/projects/edgetpu/vision/modeling/backbones/mobilenet_edgetpu.py

+# Copyright 2021 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.
+
+"""Contains definitions of mobilenet_edgetpu_v2 Networks."""
+
+# Import libraries
+
+from absl import logging
+import tensorflow as tf
+
+from official.modeling import hyperparams
+from official.projects.edgetpu.vision.modeling.mobilenet_edgetpu_v1_model import MobilenetEdgeTPU
+from official.projects.edgetpu.vision.modeling.mobilenet_edgetpu_v2_model import MobilenetEdgeTPUV2
+from official.vision.beta.modeling.backbones import factory
+
+layers = tf.keras.layers
+
+# MobileNet-EdgeTPU-V2 configs.
+MOBILENET_EDGETPU_V2_CONFIGS = frozenset([
+    'mobilenet_edgetpu_v2_tiny',
+    'mobilenet_edgetpu_v2_xs',
+    'mobilenet_edgetpu_v2_s',
+    'mobilenet_edgetpu_v2_m',
+    'mobilenet_edgetpu_v2_l',
+    'autoseg_edgetpu_backbone_xs',
+    'autoseg_edgetpu_backbone_s',
+    'autoseg_edgetpu_backbone_m',
+])
+
+# MobileNet-EdgeTPU-V1 configs.
+MOBILENET_EDGETPU_CONFIGS = frozenset([
+    'mobilenet_edgetpu',
+    'mobilenet_edgetpu_dm1p25',
+    'mobilenet_edgetpu_dm1p5',
+    'mobilenet_edgetpu_dm1p75',
+])
+
+
+def freeze_large_filters(model: tf.keras.Model, threshold: int):
+  """Freezes layer with large number of filters."""
+  for layer in model.layers:
+    if isinstance(layer.output_shape, tuple):
+      filter_size = layer.output_shape[-1]
+      if filter_size >= threshold:
+        logging.info('Freezing layer: %s', layer.name)
+        layer.trainable = False
+
+
+@factory.register_backbone_builder('mobilenet_edgetpu')
+def build_mobilenet_edgetpu(input_specs: tf.keras.layers.InputSpec,
+                            backbone_config: hyperparams.Config,
+                            **unused_kwargs) -> tf.keras.Model:
+  """Builds MobileNetEdgeTpu backbone from a config."""
+  backbone_type = backbone_config.type
+  backbone_cfg = backbone_config.get()
+  assert backbone_type == 'mobilenet_edgetpu', (f'Inconsistent backbone type '
+                                                f'{backbone_type}')
+
+  if backbone_cfg.model_id in MOBILENET_EDGETPU_V2_CONFIGS:
+    model = MobilenetEdgeTPUV2.from_name(
+        model_name=backbone_cfg.model_id,
+        overrides={
+            'batch_norm': 'tpu',
+            'rescale_input': False,
+            'resolution': input_specs.shape[1:3],
+            'backbone_only': True,
+            'features_as_dict': True,
+            'dtype': 'bfloat16'
+        },
+        model_weights_path=backbone_cfg.pretrained_checkpoint_path)
+    if backbone_cfg.freeze_large_filters:
+      freeze_large_filters(model, backbone_cfg.freeze_large_filters)
+    return model
+  elif backbone_cfg.model_id in MOBILENET_EDGETPU_CONFIGS:
+    model = MobilenetEdgeTPU.from_name(
+        model_name=backbone_cfg.model_id,
+        overrides={
+            'batch_norm': 'tpu',
+            'rescale_input': False,
+            'resolution': input_specs.shape[1:3],
+            'backbone_only': True,
+            'dtype': 'bfloat16'
+        },
+        model_weights_path=backbone_cfg.pretrained_checkpoint_path)
+    if backbone_cfg.freeze_large_filters:
+      freeze_large_filters(model, backbone_cfg.freeze_large_filters)
+    return model
+  else:
+    raise ValueError(f'Unsupported model/id type {backbone_cfg.model_id}.')

official/projects/edgetpu/vision/modeling/backbones/mobilenet_edgetpu_test.py

+# Copyright 2021 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 MobileNet."""
+
+# Import libraries
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.projects.edgetpu.vision.modeling.backbones import mobilenet_edgetpu
+
+
+class TestInputSpec:
+
+  def __init__(self, shape):
+    self.shape = shape
+
+
+class TestBackboneConfig:
+
+  def __init__(self, model_id):
+    self.model_id = model_id
+    self.freeze_large_filters = 99
+    self.pretrained_checkpoint_path = None
+    self.type = 'mobilenet_edgetpu'
+
+  def get(self):
+    return self
+
+
+class MobileNetEdgeTPUTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      ('mobilenet_edgetpu_v2_s', (1, 512, 512, 3)),
+      ('mobilenet_edgetpu_v2_l', (1, None, None, 3)),
+      ('mobilenet_edgetpu', (1, 512, 512, 3)),
+      ('mobilenet_edgetpu_dm1p25', (1, None, None, 3)),
+  )
+  def test_mobilenet_creation(self, model_id, input_shape):
+    """Test creation of MobileNet family models."""
+    tf.keras.backend.set_image_data_format('channels_last')
+
+    test_model = mobilenet_edgetpu.build_mobilenet_edgetpu(
+        input_specs=TestInputSpec(input_shape),
+        backbone_config=TestBackboneConfig(model_id))
+    self.assertGreater(len(test_model.outputs), 1)
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/edgetpu/vision/modeling/common_modules.py

+# Copyright 2021 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.
+
+"""Common modeling utilities."""
+from typing import Optional, Tuple
+# Import libraries
+import numpy as np
+import tensorflow as tf
+import tensorflow.compat.v1 as tf1
+
+from tensorflow.python.tpu import tpu_function  # pylint: disable=g-direct-tensorflow-import
+
+
+MEAN_RGB = (0.5 * 255, 0.5 * 255, 0.5 * 255)
+STDDEV_RGB = (0.5 * 255, 0.5 * 255, 0.5 * 255)
+
+
+@tf.keras.utils.register_keras_serializable(package='Vision')
+class TpuBatchNormalization(tf.keras.layers.BatchNormalization):
+  """Cross replica batch normalization."""
+
+  def __init__(self, fused: Optional[bool] = False, **kwargs):
+    if fused in (True, None):
+      raise ValueError('TpuBatchNormalization does not support fused=True.')
+    super(TpuBatchNormalization, self).__init__(fused=fused, **kwargs)
+
+  def _cross_replica_average(self, t: tf.Tensor, num_shards_per_group: int):
+    """Calculates the average value of input tensor across TPU replicas."""
+    num_shards = tpu_function.get_tpu_context().number_of_shards
+    group_assignment = None
+    if num_shards_per_group > 1:
+      if num_shards % num_shards_per_group != 0:
+        raise ValueError(
+            'num_shards: %d mod shards_per_group: %d, should be 0' %
+            (num_shards, num_shards_per_group))
+      num_groups = num_shards // num_shards_per_group
+      group_assignment = [[
+          x for x in range(num_shards) if x // num_shards_per_group == y
+      ] for y in range(num_groups)]
+    return tf1.tpu.cross_replica_sum(t, group_assignment) / tf.cast(
+        num_shards_per_group, t.dtype)
+
+  def _moments(self, inputs: tf.Tensor, reduction_axes: int, keep_dims: int):
+    """Compute the mean and variance: it overrides the original _moments."""
+    shard_mean, shard_variance = super(TpuBatchNormalization, self)._moments(
+        inputs, reduction_axes, keep_dims=keep_dims)
+
+    num_shards = tpu_function.get_tpu_context().number_of_shards or 1
+    if num_shards <= 8:  # Skip cross_replica for 2x2 or smaller slices.
+      num_shards_per_group = 1
+    else:
+      num_shards_per_group = max(8, num_shards // 8)
+    if num_shards_per_group > 1:
+      # Compute variance using: Var[X]= E[X^2] - E[X]^2.
+      shard_square_of_mean = tf.math.square(shard_mean)
+      shard_mean_of_square = shard_variance + shard_square_of_mean
+      group_mean = self._cross_replica_average(shard_mean, num_shards_per_group)
+      group_mean_of_square = self._cross_replica_average(
+          shard_mean_of_square, num_shards_per_group)
+      group_variance = group_mean_of_square - tf.math.square(group_mean)
+      return (group_mean, group_variance)
+    else:
+      return (shard_mean, shard_variance)
+
+
+def get_batch_norm(batch_norm_type: str) -> tf.keras.layers.BatchNormalization:
+  """A helper to create a batch normalization getter.
+
+  Args:
+    batch_norm_type: The type of batch normalization layer implementation. `tpu`
+     will use `TpuBatchNormalization`.
+
+  Returns:
+    An instance of `tf.keras.layers.BatchNormalization`.
+  """
+  if batch_norm_type == 'tpu':
+    return TpuBatchNormalization
+
+  return tf.keras.layers.BatchNormalization  # pytype: disable=bad-return-type  # typed-keras
+
+
+def count_params(model, trainable_only=True):
+  """Returns the count of all model parameters, or just trainable ones."""
+  if not trainable_only:
+    return model.count_params()
+  else:
+    return int(np.sum([tf.keras.backend.count_params(p)
+                       for p in model.trainable_weights]))
+
+
+def load_weights(model: tf.keras.Model,
+                 model_weights_path: str,
+                 checkpoint_format: str = 'tf_checkpoint'):
+  """Load model weights from the given file path.
+
+  Args:
+    model: the model to load weights into
+    model_weights_path: the path of the model weights
+    checkpoint_format: The source of checkpoint files. By default, we assume the
+      checkpoint is saved by tf.train.Checkpoint().save(). For legacy reasons,
+      we can also resotre checkpoint from keras model.save_weights() method by
+      setting checkpoint_format = 'keras_checkpoint'.
+  """
+  if checkpoint_format == 'tf_checkpoint':
+    checkpoint_dict = {'model': model}
+    checkpoint = tf.train.Checkpoint(**checkpoint_dict)
+    checkpoint.restore(model_weights_path).assert_existing_objects_matched()
+  elif checkpoint_format == 'keras_checkpoint':
+    # Assert makes sure load is successeful.
+    model.load_weights(model_weights_path).assert_existing_objects_matched()
+  else:
+    raise ValueError(f'Unsupported checkpoint format {checkpoint_format}.')
+
+
+def normalize_images(
+    features: tf.Tensor,
+    num_channels: int = 3,
+    dtype: str = 'float32',
+    data_format: str = 'channels_last',
+    mean_rgb: Tuple[float, ...] = MEAN_RGB,
+    stddev_rgb: Tuple[float, ...] = STDDEV_RGB,
+) -> tf.Tensor:
+  """Normalizes the input image channels with the given mean and stddev.
+
+  Args:
+    features: `Tensor` representing decoded images in float format.
+    num_channels: the number of channels in the input image tensor.
+    dtype: the dtype to convert the images to. Set to `None` to skip conversion.
+    data_format: the format of the input image tensor ['channels_first',
+      'channels_last'].
+    mean_rgb: the mean of the channels to subtract.
+    stddev_rgb: the stddev of the channels to divide.
+
+  Returns:
+    A normalized image `Tensor`.
+  """
+  if data_format == 'channels_first':
+    stats_shape = [num_channels, 1, 1]
+  else:
+    stats_shape = [1, 1, num_channels]
+
+  if dtype is not None:
+    if dtype == 'bfloat16':
+      features = tf.image.convert_image_dtype(features, dtype=tf.bfloat16)
+
+  if mean_rgb is not None:
+    mean_rgb = tf.constant(mean_rgb, shape=stats_shape, dtype=features.dtype)
+    mean_rgb = tf.broadcast_to(mean_rgb, tf.shape(features))
+    features = features - mean_rgb
+
+  if stddev_rgb is not None:
+    stddev_rgb = tf.constant(
+        stddev_rgb, shape=stats_shape, dtype=features.dtype)
+    stddev_rgb = tf.broadcast_to(stddev_rgb, tf.shape(features))
+    features = features / stddev_rgb
+
+  return features

official/projects/edgetpu/vision/modeling/heads/__init__.py

+# Copyright 2021 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.
+

official/projects/edgetpu/vision/modeling/mobilenet_edgetpu_v2_model_test.py

+# Copyright 2021 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 mobilenet_edgetpu model."""
+
+import os
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.projects.edgetpu.vision.modeling import common_modules
+from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v2_model
+
+
+class MobilenetEdgeTPUV2BuildTest(tf.test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(tf.test.TestCase, self).setUp()
+    # Ensure no model duplicates
+    tf.keras.backend.clear_session()
+
+  def test_create_mobilenet_edgetpu(self):
+    model = mobilenet_edgetpu_v2_model.MobilenetEdgeTPUV2()
+    self.assertEqual(common_modules.count_params(model), 6069657)
+
+  def test_export_tflite(self):
+    model = mobilenet_edgetpu_v2_model.MobilenetEdgeTPUV2()
+    converter = tf.lite.TFLiteConverter.from_keras_model(model)
+    converter.optimizations = [tf.lite.Optimize.DEFAULT]
+    tmp_dir = self.create_tempdir()
+    output_tflite = os.path.join(tmp_dir, 'model_quant.tflite')
+    tflite_buffer = converter.convert()
+    tf.io.gfile.GFile(output_tflite, 'wb').write(tflite_buffer)
+    self.assertTrue(tf.io.gfile.exists(output_tflite))
+
+  def test_model_save_load(self):
+    """Serializes and de-serializeds the model."""
+    model_builder = mobilenet_edgetpu_v2_model.MobilenetEdgeTPUV2
+    model = model_builder.from_name(model_name='mobilenet_edgetpu_v2')
+    # Model always has a conv2d layer followed by the input layer, and we
+    # compare the weight parameters of this layers for the original model and
+    # the save-then-load model.
+    first_conv_layer = model.get_layer('stem_conv2d')
+    kernel_tensor = first_conv_layer.trainable_weights[0].numpy()
+    model.save('/tmp/test_model')
+    loaded_model = tf.keras.models.load_model('/tmp/test_model')
+    loaded_first_conv_layer = loaded_model.get_layer('stem_conv2d')
+    loaded_kernel_tensor = loaded_first_conv_layer.trainable_weights[0].numpy()
+
+    self.assertAllClose(kernel_tensor, loaded_kernel_tensor)
+
+  def test_model_initialization_failure(self):
+    """Tests model can only be initialized with predefined model name."""
+    model_builder = mobilenet_edgetpu_v2_model.MobilenetEdgeTPUV2
+    with self.assertRaises(ValueError):
+      _ = model_builder.from_name(model_name='undefined_model_name')
+
+
+if __name__ == '__main__':
+  tf.test.main()