Create another experiment using the tfrecord path

official/projects/detr/configs/detr.py

@@ -21,10 +21,10 @@ from typing import List, Optional, Union
 from official.core import config_definitions as cfg
 from official.core import exp_factory
 from official.modeling import hyperparams
-from official.projects.detr import optimization
 from official.vision.configs import common
 from official.vision.configs import backbones
-
+from official.projects.detr import optimization
+from official.projects.detr.dataloaders import coco
 
 # pylint: disable=missing-class-docstring
 # Keep for backward compatibility.
@@ -90,11 +90,80 @@ class DetrTask(cfg.TaskConfig):
   annotation_file: Optional[str] = None
   per_category_metrics: bool = False
 
+@exp_factory.register_config_factory('detr_coco')
+def detr_coco() -> cfg.ExperimentConfig:
+  """Config to get results that matches the paper."""
+  train_batch_size = 64
+  eval_batch_size = 64
+  num_train_data = 118287
+  num_steps_per_epoch = num_train_data // train_batch_size
+  train_steps = 500 * num_steps_per_epoch  # 500 epochs
+  decay_at = train_steps - 100 * num_steps_per_epoch  # 400 epochs
+  config = cfg.ExperimentConfig(
+      task=DetrTask(
+          init_checkpoint='gs://ghpark-imagenet-tfrecord/ckpt/resnet50_imagenet',
+          init_checkpoint_modules='backbone',
+          model=Detr(
+              num_classes=81,
+              input_size=[1333, 1333, 3],
+              norm_activation=common.NormActivation(use_sync_bn=False)),
+          losses=Losses(),
+          train_data=coco.COCODataConfig(
+              file_type='tfrecord',
+              tfds_name='coco/2017',
+              tfds_split='train',
+              is_training=True,
+              global_batch_size=train_batch_size,
+              shuffle_buffer_size=1000,
+          ),
+          validation_data=coco.COCODataConfig(
+              file_type='tfrecord',
+              tfds_name='coco/2017',
+              tfds_split='validation',
+              is_training=False,
+              global_batch_size=eval_batch_size,
+              drop_remainder=False
+          )
+      ),
+      trainer=cfg.TrainerConfig(
+          train_steps=train_steps,
+          validation_steps=-1,
+          steps_per_loop=10000,
+          summary_interval=10000,
+          checkpoint_interval=10000,
+          validation_interval=10000,
+          max_to_keep=1,
+          best_checkpoint_export_subdir='best_ckpt',
+          best_checkpoint_eval_metric='AP',
+          optimizer_config=optimization.OptimizationConfig({
+              'optimizer': {
+                  'type': 'detr_adamw',
+                  'detr_adamw': {
+                      'weight_decay_rate': 1e-4,
+                      'global_clipnorm': 0.1,
+                      # Avoid AdamW legacy behavior.
+                      'gradient_clip_norm': 0.0
+                  }
+              },
+              'learning_rate': {
+                  'type': 'stepwise',
+                  'stepwise': {
+                      'boundaries': [decay_at],
+                      'values': [0.0001, 1.0e-05]
+                  }
+              },
+              })
+          ),
+      restrictions=[
+          'task.train_data.is_training != None',
+      ])
+  return config
+
 COCO_INPUT_PATH_BASE = 'gs://ghpark-tfrecords/coco'
 COCO_TRAIN_EXAMPLES = 118287
 COCO_VAL_EXAMPLES = 5000
 
-@exp_factory.register_config_factory('detr_coco')
+@exp_factory.register_config_factory('detr_coco_tfrecord')
 def detr_coco() -> cfg.ExperimentConfig:
   """Config to get results that matches the paper."""
   train_batch_size = 64

official/projects/detr/dataloaders/coco.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.
+
+"""COCO data loader for DETR."""
+
+import dataclasses
+from typing import Optional, Tuple
+import tensorflow as tf
+
+from official.core import config_definitions as cfg
+from official.core import input_reader
+from official.vision.ops import box_ops
+from official.vision.ops import preprocess_ops
+
+
+@dataclasses.dataclass
+class COCODataConfig(cfg.DataConfig):
+  """Data config for COCO."""
+  file_type: str = 'tfrecord'
+  output_size: Tuple[int, int] = (1333, 1333)
+  max_num_boxes: int = 100
+  resize_scales: Tuple[int, ...] = (
+      480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
+
+
+class COCODataLoader():
+  """A class to load dataset for COCO detection task."""
+
+  def __init__(self, params: COCODataConfig):
+    self._params = params
+
+  def preprocess(self, inputs):
+    """Preprocess COCO for DETR."""
+    image = inputs['image']
+    boxes = inputs['objects']['bbox']
+    classes = inputs['objects']['label'] + 1
+    is_crowd = inputs['objects']['is_crowd']
+
+    image = preprocess_ops.normalize_image(image)
+    if self._params.is_training:
+      image, boxes, _ = preprocess_ops.random_horizontal_flip(image, boxes)
+
+      do_crop = tf.greater(tf.random.uniform([]), 0.5)
+      if do_crop:
+        # Rescale
+        boxes = box_ops.denormalize_boxes(boxes, tf.shape(image)[:2])
+        index = tf.random.categorical(tf.zeros([1, 3]), 1)[0]
+        scales = tf.gather([400.0, 500.0, 600.0], index, axis=0)
+        short_side = scales[0]
+        image, image_info = preprocess_ops.resize_image(image, short_side)
+        boxes = preprocess_ops.resize_and_crop_boxes(boxes,
+                                                     image_info[2, :],
+                                                     image_info[1, :],
+                                                     image_info[3, :])
+        boxes = box_ops.normalize_boxes(boxes, image_info[1, :])
+
+        # Do croping
+        shape = tf.cast(image_info[1], dtype=tf.int32)
+        h = tf.random.uniform(
+            [], 384, tf.math.minimum(shape[0], 600), dtype=tf.int32)
+        w = tf.random.uniform(
+            [], 384, tf.math.minimum(shape[1], 600), dtype=tf.int32)
+        i = tf.random.uniform([], 0, shape[0] - h + 1, dtype=tf.int32)
+        j = tf.random.uniform([], 0, shape[1] - w + 1, dtype=tf.int32)
+        image = tf.image.crop_to_bounding_box(image, i, j, h, w)
+        boxes = tf.clip_by_value(
+            (boxes[..., :] * tf.cast(
+                tf.stack([shape[0], shape[1], shape[0], shape[1]]),
+                dtype=tf.float32) -
+             tf.cast(tf.stack([i, j, i, j]), dtype=tf.float32)) /
+            tf.cast(tf.stack([h, w, h, w]), dtype=tf.float32), 0.0, 1.0)
+      scales = tf.constant(
+          self._params.resize_scales,
+          dtype=tf.float32)
+      index = tf.random.categorical(tf.zeros([1, 11]), 1)[0]
+      scales = tf.gather(scales, index, axis=0)
+    else:
+      scales = tf.constant([self._params.resize_scales[-1]], tf.float32)
+
+    image_shape = tf.shape(image)[:2]
+    boxes = box_ops.denormalize_boxes(boxes, image_shape)
+    gt_boxes = boxes
+    short_side = scales[0]
+    image, image_info = preprocess_ops.resize_image(
+        image,
+        short_side,
+        max(self._params.output_size))
+    boxes = preprocess_ops.resize_and_crop_boxes(boxes,
+                                                 image_info[2, :],
+                                                 image_info[1, :],
+                                                 image_info[3, :])
+    boxes = box_ops.normalize_boxes(boxes, image_info[1, :])
+
+    # Filters out ground truth boxes that are all zeros.
+    indices = box_ops.get_non_empty_box_indices(boxes)
+    boxes = tf.gather(boxes, indices)
+    classes = tf.gather(classes, indices)
+    is_crowd = tf.gather(is_crowd, indices)
+    boxes = box_ops.yxyx_to_cycxhw(boxes)
+
+    image = tf.image.pad_to_bounding_box(
+        image, 0, 0, self._params.output_size[0], self._params.output_size[1])
+    labels = {
+        'classes':
+            preprocess_ops.clip_or_pad_to_fixed_size(
+                classes, self._params.max_num_boxes),
+        'boxes':
+            preprocess_ops.clip_or_pad_to_fixed_size(
+                boxes, self._params.max_num_boxes)
+    }
+    if not self._params.is_training:
+      labels.update({
+          'id':
+              inputs['image/id'],
+          'image_info':
+              image_info,
+          'is_crowd':
+              preprocess_ops.clip_or_pad_to_fixed_size(
+                  is_crowd, self._params.max_num_boxes),
+          'gt_boxes':
+              preprocess_ops.clip_or_pad_to_fixed_size(
+                  gt_boxes, self._params.max_num_boxes),
+      })
+
+    return image, labels
+
+  def _transform_and_batch_fn(
+      self,
+      dataset,
+      input_context: Optional[tf.distribute.InputContext] = None):
+    """Preprocess and batch."""
+    dataset = dataset.map(
+        self.preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+    per_replica_batch_size = input_context.get_per_replica_batch_size(
+        self._params.global_batch_size
+    ) if input_context else self._params.global_batch_size
+    dataset = dataset.batch(
+        per_replica_batch_size, drop_remainder=self._params.is_training)
+    return dataset
+
+  def load(self, input_context: Optional[tf.distribute.InputContext] = None):
+    """Returns a tf.dataset.Dataset."""
+    reader = input_reader.InputReader(
+        params=self._params,
+        decoder_fn=None,
+        transform_and_batch_fn=self._transform_and_batch_fn)
+    return reader.read(input_context)

official/projects/detr/dataloaders/coco_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 tensorflow_models.official.projects.detr.dataloaders.coco."""
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+from official.projects.detr.dataloaders import coco
+
+
+def _gen_fn():
+  h = np.random.randint(0, 300)
+  w = np.random.randint(0, 300)
+  num_boxes = np.random.randint(0, 50)
+  return {
+      'image': np.ones(shape=(h, w, 3), dtype=np.uint8),
+      'image/id': np.random.randint(0, 100),
+      'image/filename': 'test',
+      'objects': {
+          'is_crowd': np.ones(shape=(num_boxes), dtype=np.bool),
+          'bbox': np.ones(shape=(num_boxes, 4), dtype=np.float32),
+          'label': np.ones(shape=(num_boxes), dtype=np.int64),
+          'id': np.ones(shape=(num_boxes), dtype=np.int64),
+          'area': np.ones(shape=(num_boxes), dtype=np.int64),
+      }
+  }
+
+
+class CocoDataloaderTest(tf.test.TestCase, parameterized.TestCase):
+
+  def test_load_dataset(self):
+    output_size = 1280
+    max_num_boxes = 100
+    batch_size = 2
+    data_config = coco.COCODataConfig(
+        tfds_name='coco/2017',
+        tfds_split='validation',
+        is_training=False,
+        global_batch_size=batch_size,
+        output_size=(output_size, output_size),
+        max_num_boxes=max_num_boxes,
+        )
+
+    num_examples = 10
+    def as_dataset(self, *args, **kwargs):
+      del args
+      del kwargs
+      return tf.data.Dataset.from_generator(
+          lambda: (_gen_fn() for i in range(num_examples)),
+          output_types=self.info.features.dtype,
+          output_shapes=self.info.features.shape,
+      )
+
+    with tfds.testing.mock_data(num_examples=num_examples,
+                                as_dataset_fn=as_dataset):
+      dataset = coco.COCODataLoader(data_config).load()
+      dataset_iter = iter(dataset)
+      images, labels = next(dataset_iter)
+      self.assertEqual(images.shape, (batch_size, output_size, output_size, 3))
+      self.assertEqual(labels['classes'].shape, (batch_size, max_num_boxes))
+      self.assertEqual(labels['boxes'].shape, (batch_size, max_num_boxes, 4))
+      self.assertEqual(labels['id'].shape, (batch_size,))
+      self.assertEqual(
+          labels['image_info'].shape, (batch_size, 4, 2))
+      self.assertEqual(labels['is_crowd'].shape, (batch_size, max_num_boxes))
+
+  @parameterized.named_parameters(
+      ('training', True),
+      ('validation', False))
+  def test_preprocess(self, is_training):
+    output_size = 1280
+    max_num_boxes = 100
+    batch_size = 2
+    data_config = coco.COCODataConfig(
+        tfds_name='coco/2017',
+        tfds_split='validation',
+        is_training=is_training,
+        global_batch_size=batch_size,
+        output_size=(output_size, output_size),
+        max_num_boxes=max_num_boxes,
+        )
+
+    dl = coco.COCODataLoader(data_config)
+    inputs = _gen_fn()
+    image, label = dl.preprocess(inputs)
+    self.assertEqual(image.shape, (output_size, output_size, 3))
+    self.assertEqual(label['classes'].shape, (max_num_boxes))
+    self.assertEqual(label['boxes'].shape, (max_num_boxes, 4))
+    if not is_training:
+      self.assertDTypeEqual(label['id'], int)
+      self.assertEqual(
+          label['image_info'].shape, (4, 2))
+      self.assertEqual(label['is_crowd'].shape, (max_num_boxes))
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/detr/experiments/detr_r50_300epochs.sh

@@ -3,4 +3,4 @@ python3 official/projects/detr/train.py \
   --experiment=detr_coco \
   --mode=train_and_eval \
   --model_dir=/tmp/logging_dir/ \
-  --params_override=task.init_ckpt='gs://tf_model_garden/vision/resnet50_imagenet/ckpt-62400',trainer.train_steps=554400
+  --params_override=task.init_checkpoint='gs://tf_model_garden/vision/resnet50_imagenet/ckpt-62400',trainer.train_steps=554400

official/projects/detr/experiments/detr_r50_500epochs.sh

@@ -3,4 +3,4 @@ python3 official/projects/detr/train.py \
   --experiment=detr_coco \
   --mode=train_and_eval \
   --model_dir=/tmp/logging_dir/ \
-  --params_override=task.init_ckpt='gs://tf_model_garden/vision/resnet50_imagenet/ckpt-62400'
+  --params_override=task.init_checkpoint='gs://tf_model_garden/vision/resnet50_imagenet/ckpt-62400'