Internal change

PiperOrigin-RevId: 356306227

official/vision/beta/configs/experiments/retinanet/resnet50fpn_coco_tfds_tpu.yaml

+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'bfloat16'
+task:
+  annotation_file: ''  # Can't use annotation file when tfds is used.
+  losses:
+    l2_weight_decay: 0.0001
+  model:
+    num_classes: 91
+    max_level: 7
+    min_level: 3
+    input_size: [640, 640, 3]
+    norm_activation:
+      activation: relu
+      norm_epsilon: 0.001
+      norm_momentum: 0.99
+      use_sync_bn: true
+  train_data:
+    tfds_name: 'coco/2017'
+    tfds_split: 'train'
+    drop_remainder: true
+    dtype: bfloat16
+    global_batch_size: 256
+    input_path: ''
+    is_training: true
+    shuffle_buffer_size: 1000
+  validation_data:
+    tfds_name: 'coco/2017'
+    tfds_split: 'validation'
+    drop_remainder: true
+    dtype: bfloat16
+    global_batch_size: 8
+    input_path: ''
+    is_training: false

official/vision/beta/configs/experiments/semantic_segmentation/deeplabv3plus_resnet101_cityscapes_tfds_tpu.yaml

+# Use your own cityscapes preprocessed dataset. 79% meanIoU.
+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'float32'
+task:
+  model:
+    num_classes: 19
+    input_size: [null, null, 3]
+    backbone:
+      type: 'dilated_resnet'
+      dilated_resnet:
+        model_id: 101
+        output_stride: 16
+        stem_type: 'v1'
+        se_ratio: 0.25
+        stochastic_depth_drop_rate: 0.2
+        multigrid: [1, 2, 4]
+        last_stage_repeats: 1
+    decoder:
+      aspp:
+        pool_kernel_size: [512, 1024]
+    head:
+      feature_fusion: 'deeplabv3plus'
+      low_level: 2
+      low_level_num_filters: 48
+    norm_activation:
+      activation: 'swish'
+      norm_epsilon: 0.001
+      norm_momentum: 0.99
+      use_sync_bn: true
+  losses:
+    top_k_percent_pixels: 1.0  # only backpropagate loss for the topk 100% pixels.
+  train_data:
+    output_size: [512, 1024]
+    train_on_crops: true
+    input_path: ''
+    tfds_name: 'cityscapes/semantic_segmentation'
+    tfds_split: 'train'
+    is_training: true
+    global_batch_size: 16
+    dtype: 'float32'
+    aug_rand_hflip: true
+    aug_scale_max: 2.0
+    aug_scale_min: 0.5
+  validation_data:
+    output_size: [1024, 2048]
+    input_path: ''
+    tfds_name: 'cityscapes/semantic_segmentation'
+    tfds_split: 'validation'
+    is_training: false
+    global_batch_size: 16
+    dtype: 'float32'
+    drop_remainder: false
+    resize_eval_groundtruth: true
+trainer:
+  optimizer_config:
+    learning_rate:
+      polynomial:
+        decay_steps: 90000
+        initial_learning_rate: 0.01
+        power: 0.9
+      type: polynomial
+    optimizer:
+      sgd:
+        momentum: 0.9
+      type: sgd
+    warmup:
+      linear:
+        name: linear
+        warmup_learning_rate: 0
+        warmup_steps: 925
+      type: linear
+  steps_per_loop: 185
+  summary_interval: 185
+  train_steps: 90000
+  validation_interval: 185
+  validation_steps: 31
+  checkpoint_interval: 185

official/vision/beta/dataloaders/tfds_detection_decoders.py

+# Copyright 2020 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.
+# ==============================================================================
+"""TFDS detection decoders."""
+
+import tensorflow as tf
+from official.vision.beta.dataloaders import decoder
+
+
+class MSCOCODecoder(decoder.Decoder):
+  """A tf.Example decoder for tfds coco datasets."""
+
+  def decode(self, serialized_example):
+    """Decode the serialized example.
+
+    Args:
+      serialized_example: a dictonary example produced by tfds.
+
+    Returns:
+      decoded_tensors: a dictionary of tensors with the following fields:
+        - source_id: a string scalar tensor.
+        - image: a uint8 tensor of shape [None, None, 3].
+        - height: an integer scalar tensor.
+        - width: an integer scalar tensor.
+        - groundtruth_classes: a int64 tensor of shape [None].
+        - groundtruth_is_crowd: a bool tensor of shape [None].
+        - groundtruth_area: a float32 tensor of shape [None].
+        - groundtruth_boxes: a float32 tensor of shape [None, 4].
+    """
+
+    decoded_tensors = {
+        'source_id': tf.strings.as_string(serialized_example['image/id']),
+        'image': serialized_example['image'],
+        'height': tf.cast(tf.shape(serialized_example['image'])[0], tf.int64),
+        'width': tf.cast(tf.shape(serialized_example['image'])[1], tf.int64),
+        'groundtruth_classes': serialized_example['objects']['label'],
+        'groundtruth_is_crowd': serialized_example['objects']['is_crowd'],
+        'groundtruth_area': tf.cast(
+            serialized_example['objects']['area'], tf.float32),
+        'groundtruth_boxes': serialized_example['objects']['bbox'],
+    }
+    return decoded_tensors
+
+
+TFDS_ID_TO_DECODER_MAP = {
+    'coco/2017': MSCOCODecoder,
+    'coco/2014': MSCOCODecoder,
+    'coco': MSCOCODecoder
+}

official/vision/beta/dataloaders/tfds_segmentation_decoders.py

+# Copyright 2020 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.
+# ==============================================================================
+"""TFDS Semantic Segmentation decoders."""
+
+import tensorflow as tf
+from official.vision.beta.dataloaders import decoder
+
+
+class CityScapesDecorder(decoder.Decoder):
+  """A tf.Example decoder for tfds cityscapes datasets."""
+
+  def __init__(self):
+    # Original labels to trainable labels map, 255 is the ignore class.
+    self._label_map = {
+        -1: 255,
+        0: 255,
+        1: 255,
+        2: 255,
+        3: 255,
+        4: 255,
+        5: 255,
+        6: 255,
+        7: 0,
+        8: 1,
+        9: 255,
+        10: 255,
+        11: 2,
+        12: 3,
+        13: 4,
+        14: 255,
+        15: 255,
+        16: 255,
+        17: 5,
+        18: 255,
+        19: 6,
+        20: 7,
+        21: 8,
+        22: 9,
+        23: 10,
+        24: 11,
+        25: 12,
+        26: 13,
+        27: 14,
+        28: 15,
+        29: 255,
+        30: 255,
+        31: 16,
+        32: 17,
+        33: 18,
+    }
+
+  def decode(self, serialized_example):
+    # Convert labels according to the self._label_map
+    label = serialized_example['segmentation_label']
+    for original_label in self._label_map:
+      label = tf.where(label == original_label,
+                       self._label_map[original_label] * tf.ones_like(label),
+                       label)
+    sample_dict = {
+        'image/encoded':
+            tf.io.encode_jpeg(serialized_example['image_left'], quality=100),
+        'image/height': serialized_example['image_left'].shape[0],
+        'image/width': serialized_example['image_left'].shape[1],
+        'image/segmentation/class/encoded':
+            tf.io.encode_png(label),
+    }
+    return sample_dict
+
+
+TFDS_ID_TO_DECODER_MAP = {
+    'cityscapes': CityScapesDecorder,
+    'cityscapes/semantic_segmentation': CityScapesDecorder,
+    'cityscapes/semantic_segmentation_extra': CityScapesDecorder,
+}

official/vision/beta/tasks/retinanet.py

@@ -25,6 +25,7 @@ from official.vision import keras_cv
 from official.vision.beta.configs import retinanet as exp_cfg
 from official.vision.beta.dataloaders import retinanet_input
 from official.vision.beta.dataloaders import tf_example_decoder
+from official.vision.beta.dataloaders import tfds_detection_decoders
 from official.vision.beta.dataloaders import tf_example_label_map_decoder
 from official.vision.beta.evaluation import coco_evaluator
 from official.vision.beta.modeling import factory
@@ -84,16 +85,25 @@ class RetinaNetTask(base_task.Task):
 
   def build_inputs(self, params, input_context=None):
     """Build input dataset."""
-    decoder_cfg = params.decoder.get()
-    if params.decoder.type == 'simple_decoder':
-      decoder = tf_example_decoder.TfExampleDecoder(
-          regenerate_source_id=decoder_cfg.regenerate_source_id)
-    elif params.decoder.type == 'label_map_decoder':
-      decoder = tf_example_label_map_decoder.TfExampleDecoderLabelMap(
-          label_map=decoder_cfg.label_map,
-          regenerate_source_id=decoder_cfg.regenerate_source_id)
+
+    if params.tfds_name:
+      if params.tfds_name in tfds_detection_decoders.TFDS_ID_TO_DECODER_MAP:
+        decoder = tfds_detection_decoders.TFDS_ID_TO_DECODER_MAP[
+            params.tfds_name]()
+      else:
+        raise ValueError('TFDS {} is not supported'.format(params.tfds_name))
     else:
-      raise ValueError('Unknown decoder type: {}!'.format(params.decoder.type))
+      decoder_cfg = params.decoder.get()
+      if params.decoder.type == 'simple_decoder':
+        decoder = tf_example_decoder.TfExampleDecoder(
+            regenerate_source_id=decoder_cfg.regenerate_source_id)
+      elif params.decoder.type == 'label_map_decoder':
+        decoder = tf_example_label_map_decoder.TfExampleDecoderLabelMap(
+            label_map=decoder_cfg.label_map,
+            regenerate_source_id=decoder_cfg.regenerate_source_id)
+      else:
+        raise ValueError('Unknown decoder type: {}!'.format(
+            params.decoder.type))
 
     parser = retinanet_input.Parser(
         output_size=self.task_config.model.input_size[:2],
@@ -169,10 +179,13 @@ class RetinaNetTask(base_task.Task):
       metrics.append(tf.keras.metrics.Mean(name, dtype=tf.float32))
 
     if not training:
+      if self.task_config.validation_data.tfds_name and self.task_config.annotation_file:
+        raise ValueError(
+            "Can't evaluate using annotation file when TFDS is used.")
       self.coco_metric = coco_evaluator.COCOEvaluator(
-          annotation_file=self._task_config.annotation_file,
+          annotation_file=self.task_config.annotation_file,
           include_mask=False,
-          per_category_metrics=self._task_config.per_category_metrics)
+          per_category_metrics=self.task_config.per_category_metrics)
 
     return metrics
 

official/vision/beta/tasks/semantic_segmentation.py

@@ -23,6 +23,7 @@ from official.core import input_reader
 from official.core import task_factory
 from official.vision.beta.configs import semantic_segmentation as exp_cfg
 from official.vision.beta.dataloaders import segmentation_input
+from official.vision.beta.dataloaders import tfds_segmentation_decoders
 from official.vision.beta.evaluation import segmentation_metrics
 from official.vision.beta.losses import segmentation_losses
 from official.vision.beta.modeling import factory
@@ -83,7 +84,15 @@ class SemanticSegmentationTask(base_task.Task):
 
     ignore_label = self.task_config.losses.ignore_label
 
-    decoder = segmentation_input.Decoder()
+    if params.tfds_name:
+      if params.tfds_name in tfds_segmentation_decoders.TFDS_ID_TO_DECODER_MAP:
+        decoder = tfds_segmentation_decoders.TFDS_ID_TO_DECODER_MAP[
+            params.tfds_name]()
+      else:
+        raise ValueError('TFDS {} is not supported'.format(params.tfds_name))
+    else:
+      decoder = segmentation_input.Decoder()
+
     parser = segmentation_input.Parser(
         output_size=params.output_size,
         train_on_crops=params.train_on_crops,