tf-models-nightly 2.19.0.dev20250108__py2.py3-none-any.whl → 2.19.0.dev20250109__py2.py3-none-any.whl

official/projects/detr/optimization.py

@@ -0,0 +1,151 @@
+# Copyright 2024 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.
+
+"""Customized optimizer to match paper results."""
+
+import dataclasses
+import tensorflow as tf, tf_keras
+from official.modeling import optimization
+from official.nlp import optimization as nlp_optimization
+
+
+@dataclasses.dataclass
+class DETRAdamWConfig(optimization.AdamWeightDecayConfig):
+  pass
+
+
+@dataclasses.dataclass
+class OptimizerConfig(optimization.OptimizerConfig):
+  detr_adamw: DETRAdamWConfig = dataclasses.field(
+      default_factory=DETRAdamWConfig
+  )
+
+
+@dataclasses.dataclass
+class OptimizationConfig(optimization.OptimizationConfig):
+  """Configuration for optimizer and learning rate schedule.
+
+  Attributes:
+    optimizer: optimizer oneof config.
+    ema: optional exponential moving average optimizer config, if specified, ema
+      optimizer will be used.
+    learning_rate: learning rate oneof config.
+    warmup: warmup oneof config.
+  """
+  optimizer: OptimizerConfig = dataclasses.field(
+      default_factory=OptimizerConfig
+  )
+
+
+# TODO(frederickliu): figure out how to make this configuable.
+# TODO(frederickliu): Study if this is needed.
+class _DETRAdamW(nlp_optimization.AdamWeightDecay):
+  """Custom AdamW to support different lr scaling for backbone.
+
+  The code is copied from AdamWeightDecay and Adam with learning scaling.
+  """
+
+  def _resource_apply_dense(self, grad, var, apply_state=None):
+    lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
+    apply_state = kwargs['apply_state']
+    if 'detr' not in var.name:
+      lr_t *= 0.1
+    decay = self._decay_weights_op(var, lr_t, apply_state)
+    with tf.control_dependencies([decay]):
+      var_device, var_dtype = var.device, var.dtype.base_dtype
+      coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                      or self._fallback_apply_state(var_device, var_dtype))
+
+      m = self.get_slot(var, 'm')
+      v = self.get_slot(var, 'v')
+      lr = coefficients[
+          'lr_t'] * 0.1 if 'detr' not in var.name else coefficients['lr_t']
+
+      if not self.amsgrad:
+        return tf.raw_ops.ResourceApplyAdam(
+            var=var.handle,
+            m=m.handle,
+            v=v.handle,
+            beta1_power=coefficients['beta_1_power'],
+            beta2_power=coefficients['beta_2_power'],
+            lr=lr,
+            beta1=coefficients['beta_1_t'],
+            beta2=coefficients['beta_2_t'],
+            epsilon=coefficients['epsilon'],
+            grad=grad,
+            use_locking=self._use_locking)
+      else:
+        vhat = self.get_slot(var, 'vhat')
+        return tf.raw_ops.ResourceApplyAdamWithAmsgrad(
+            var=var.handle,
+            m=m.handle,
+            v=v.handle,
+            vhat=vhat.handle,
+            beta1_power=coefficients['beta_1_power'],
+            beta2_power=coefficients['beta_2_power'],
+            lr=lr,
+            beta1=coefficients['beta_1_t'],
+            beta2=coefficients['beta_2_t'],
+            epsilon=coefficients['epsilon'],
+            grad=grad,
+            use_locking=self._use_locking)
+
+  def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
+    lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state)
+    apply_state = kwargs['apply_state']
+    if 'detr' not in var.name:
+      lr_t *= 0.1
+    decay = self._decay_weights_op(var, lr_t, apply_state)
+    with tf.control_dependencies([decay]):
+      var_device, var_dtype = var.device, var.dtype.base_dtype
+      coefficients = ((apply_state or {}).get((var_device, var_dtype))
+                      or self._fallback_apply_state(var_device, var_dtype))
+
+      # m_t = beta1 * m + (1 - beta1) * g_t
+      m = self.get_slot(var, 'm')
+      m_scaled_g_values = grad * coefficients['one_minus_beta_1_t']
+      m_t = tf.compat.v1.assign(m, m * coefficients['beta_1_t'],
+                                use_locking=self._use_locking)
+      with tf.control_dependencies([m_t]):
+        m_t = self._resource_scatter_add(m, indices, m_scaled_g_values)
+
+      # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+      v = self.get_slot(var, 'v')
+      v_scaled_g_values = (grad * grad) * coefficients['one_minus_beta_2_t']
+      v_t = tf.compat.v1.assign(v, v * coefficients['beta_2_t'],
+                                use_locking=self._use_locking)
+      with tf.control_dependencies([v_t]):
+        v_t = self._resource_scatter_add(v, indices, v_scaled_g_values)
+      lr = coefficients[
+          'lr_t'] * 0.1 if 'detr' not in var.name else coefficients['lr_t']
+      if not self.amsgrad:
+        v_sqrt = tf.sqrt(v_t)
+        var_update = tf.compat.v1.assign_sub(
+            var, lr * m_t / (v_sqrt + coefficients['epsilon']),
+            use_locking=self._use_locking)
+        return tf.group(*[var_update, m_t, v_t])
+      else:
+        v_hat = self.get_slot(var, 'vhat')
+        v_hat_t = tf.maximum(v_hat, v_t)
+        with tf.control_dependencies([v_hat_t]):
+          v_hat_t = tf.compat.v1.assign(
+              v_hat, v_hat_t, use_locking=self._use_locking)
+        v_hat_sqrt = tf.sqrt(v_hat_t)
+        var_update = tf.compat.v1.assign_sub(
+            var,
+            lr* m_t / (v_hat_sqrt + coefficients['epsilon']),
+            use_locking=self._use_locking)
+        return tf.group(*[var_update, m_t, v_t, v_hat_t])
+
+optimization.register_optimizer_cls('detr_adamw', _DETRAdamW)

official/projects/detr/serving/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 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/detr/serving/export_module.py

@@ -0,0 +1,103 @@
+# Copyright 2024 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.
+
+"""Export module for DETR model."""
+import tensorflow as tf, tf_keras
+
+from official.projects.detr.modeling import detr
+from official.vision.modeling import backbones
+from official.vision.ops import preprocess_ops
+from official.vision.serving import detection
+
+
+class DETRModule(detection.DetectionModule):
+  """DETR detection module."""
+
+  def _build_model(self) -> tf_keras.Model:
+    input_specs = tf_keras.layers.InputSpec(shape=[self._batch_size] +
+                                            self._input_image_size +
+                                            [self._num_channels])
+
+    backbone = backbones.factory.build_backbone(
+        input_specs=input_specs,
+        backbone_config=self.params.task.model.backbone,
+        norm_activation_config=self.params.task.model.norm_activation)
+
+    model = detr.DETR(backbone, self.params.task.model.backbone_endpoint_name,
+                      self.params.task.model.num_queries,
+                      self.params.task.model.hidden_size,
+                      self.params.task.model.num_classes,
+                      self.params.task.model.num_encoder_layers,
+                      self.params.task.model.num_decoder_layers)
+    model(tf_keras.Input(input_specs.shape[1:]))
+    return model
+
+  def _build_inputs(self, image: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]:
+    """Builds detection model inputs for serving."""
+    # Normalizes image with mean and std pixel values.
+    image = preprocess_ops.normalize_image(
+        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
+
+    image, image_info = preprocess_ops.resize_image(
+        image, size=self._input_image_size)
+
+    return image, image_info
+
+  def serve(self, images: tf.Tensor) -> dict[str, tf.Tensor]:
+    """Cast image to float and run inference.
+
+    Args:
+      images: uint8 Tensor of shape [batch_size, None, None, 3]
+
+    Returns:
+      Tensor holding classification output logits.
+    """
+    # Skip image preprocessing when input_type is tflite so it is compatible
+    # with TFLite quantization.
+    image_info = None
+    if self._input_type != 'tflite':
+      with tf.device('cpu:0'):
+        images = tf.cast(images, dtype=tf.float32)
+
+        images_spec = tf.TensorSpec(
+            shape=self._input_image_size + [3], dtype=tf.float32)
+        image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
+
+        images, image_info = tf.nest.map_structure(
+            tf.identity,
+            tf.map_fn(
+                self._build_inputs,
+                elems=images,
+                fn_output_signature=(images_spec, image_info_spec),
+                parallel_iterations=32))
+
+    outputs = self.inference_step(images)[-1]
+    outputs = {
+        'detection_boxes': outputs['detection_boxes'],
+        'detection_scores': outputs['detection_scores'],
+        'detection_classes': outputs['detection_classes'],
+        'num_detections': outputs['num_detections']
+    }
+    if image_info is not None:
+      outputs['detection_boxes'] = outputs['detection_boxes'] * tf.expand_dims(
+          tf.concat([
+              image_info[:, 1:2, 0], image_info[:, 1:2, 1],
+              image_info[:, 1:2, 0], image_info[:, 1:2, 1]
+          ],
+                    axis=1),
+          axis=1)
+
+      outputs.update({'image_info': image_info})
+
+    return outputs

official/projects/detr/serving/export_module_test.py

@@ -0,0 +1,98 @@
+# Copyright 2024 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.
+
+"""Test for DETR export module."""
+
+import io
+import os
+
+from absl.testing import parameterized
+import numpy as np
+from PIL import Image
+import tensorflow as tf, tf_keras
+
+from official.core import exp_factory
+from official.projects.detr.configs import detr as exp_cfg  # pylint: disable=unused-import
+from official.projects.detr.serving import export_module
+
+
+class ExportModuleTest(tf.test.TestCase, parameterized.TestCase):
+
+  def _get_module(self, input_type):
+    params = exp_factory.get_exp_config('detr_coco')
+    return export_module.DETRModule(
+        params,
+        batch_size=1,
+        input_image_size=[384, 384],
+        input_type=input_type)
+
+  def _export_from_module(self, module, input_type, save_directory):
+    signatures = module.get_inference_signatures(
+        {input_type: 'serving_default'})
+    tf.saved_model.save(module, save_directory, signatures=signatures)
+
+  def _get_dummy_input(self, input_type):
+    """Gets dummy input for the given input type."""
+
+    if input_type == 'image_tensor':
+      return tf.zeros((1, 384, 384, 3), dtype=np.uint8)
+    elif input_type == 'image_bytes':
+      image = Image.fromarray(np.zeros((384, 384, 3), dtype=np.uint8))
+      byte_io = io.BytesIO()
+      image.save(byte_io, 'PNG')
+      return [byte_io.getvalue()]
+    elif input_type == 'tf_example':
+      image_tensor = tf.zeros((384, 384, 3), dtype=tf.uint8)
+      encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
+      example = tf.train.Example(
+          features=tf.train.Features(
+              feature={
+                  'image/encoded':
+                      tf.train.Feature(
+                          bytes_list=tf.train.BytesList(value=[encoded_jpeg])),
+              })).SerializeToString()
+      return [example]
+
+  @parameterized.parameters(
+      {'input_type': 'image_tensor'},
+      {'input_type': 'image_bytes'},
+      {'input_type': 'tf_example'},
+  )
+  def test_export(self, input_type='image_tensor'):
+    tmp_dir = self.get_temp_dir()
+    module = self._get_module(input_type)
+    self._export_from_module(module, input_type, tmp_dir)
+
+    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
+    self.assertTrue(
+        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index')))
+    self.assertTrue(
+        os.path.exists(
+            os.path.join(tmp_dir, 'variables',
+                         'variables.data-00000-of-00001')))
+
+    imported = tf.saved_model.load(tmp_dir)
+    predict_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type)
+    outputs = predict_fn(tf.constant(images))
+
+    self.assertNotEmpty(outputs['detection_boxes'])
+    self.assertNotEmpty(outputs['detection_classes'])
+    self.assertNotEmpty(outputs['detection_scores'])
+    self.assertNotEmpty(outputs['num_detections'])
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/detr/serving/export_saved_model.py

@@ -0,0 +1,109 @@
+# Copyright 2024 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.
+
+r"""Vision models export binary for serving/inference.
+
+To export a trained checkpoint in saved_model format (shell script):
+
+EXPERIMENT_TYPE = XX
+CHECKPOINT_PATH = XX
+EXPORT_DIR_PATH = XX
+export_saved_model --experiment=${EXPERIMENT_TYPE} \
+                   --export_dir=${EXPORT_DIR_PATH}/ \
+                   --checkpoint_path=${CHECKPOINT_PATH} \
+                   --batch_size=2 \
+                   --input_image_size=224,224
+
+To serve (python):
+
+export_dir_path = XX
+input_type = XX
+input_images = XX
+imported = tf.saved_model.load(export_dir_path)
+model_fn = imported.signatures['serving_default']
+output = model_fn(input_images)
+"""
+
+from absl import app
+from absl import flags
+
+from official.core import exp_factory
+from official.modeling import hyperparams
+from official.projects.detr.configs import detr as exp_cfg  # pylint: disable=unused-import
+from official.projects.detr.serving import export_module
+from official.vision.serving import export_saved_model_lib
+
+FLAGS = flags.FLAGS
+
+_EXPERIMENT = flags.DEFINE_string('experiment', None,
+                                  'experiment type, e.g. detr_coco')
+_EXPORT_DIR = flags.DEFINE_string('export_dir', None, 'The export directory.')
+_CHECKPOINT_PATH = flags.DEFINE_string('checkpoint_path', None,
+                                       'Checkpoint path.')
+_CONFIG_FILE = flags.DEFINE_multi_string(
+    'config_file',
+    default=None,
+    help='YAML/JSON files which specifies overrides. The override order '
+    'follows the order of args. Note that each file '
+    'can be used as an override template to override the default parameters '
+    'specified in Python. If the same parameter is specified in both '
+    '`--config_file` and `--params_override`, `config_file` will be used '
+    'first, followed by params_override.')
+_PARAMS_OVERRIDE = flags.DEFINE_string(
+    'params_override', '',
+    'The JSON/YAML file or string which specifies the parameter to be overriden'
+    ' on top of `config_file` template.')
+_BATCH_SIZE = flags.DEFINE_integer('batch_size', None, 'The batch size.')
+_IMAGE_TYPE = flags.DEFINE_string(
+    'input_type', 'image_tensor',
+    'One of `image_tensor`, `image_bytes`, `tf_example` and `tflite`.')
+_INPUT_IMAGE_SIZE = flags.DEFINE_string(
+    'input_image_size', '224,224',
+    'The comma-separated string of two integers representing the height,width '
+    'of the input to the model.')
+
+
+def main(_):
+
+  params = exp_factory.get_exp_config(_EXPERIMENT.value)
+  for config_file in _CONFIG_FILE.value or []:
+    params = hyperparams.override_params_dict(
+        params, config_file, is_strict=False)
+  if _PARAMS_OVERRIDE.value:
+    params = hyperparams.override_params_dict(
+        params, _PARAMS_OVERRIDE.value, is_strict=False)
+
+  params.validate()
+  params.lock()
+
+  input_image_size = [int(x) for x in _INPUT_IMAGE_SIZE.value.split(',')]
+  module = export_module.DETRModule(
+      params=params,
+      batch_size=_BATCH_SIZE.value,
+      input_image_size=input_image_size,
+      input_type=_IMAGE_TYPE.value,
+      num_channels=3)
+
+  export_saved_model_lib.export_inference_graph(
+      input_type=_IMAGE_TYPE.value,
+      batch_size=_BATCH_SIZE.value,
+      input_image_size=input_image_size,
+      params=params,
+      checkpoint_path=_CHECKPOINT_PATH.value,
+      export_dir=_EXPORT_DIR.value,
+      export_module=module)
+
+
+if __name__ == '__main__':
+  app.run(main)

official/projects/detr/tasks/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 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.
+