Merge pull request #10338 from srihari-humbarwadi:readme

PiperOrigin-RevId: 413033276

official/vision/beta/projects/centernet/modeling/layers/cn_nn_blocks.py

@@ -123,7 +123,6 @@ def _make_repeated_residual_blocks(
   return tf.keras.Sequential(blocks)
 
 
-@tf.keras.utils.register_keras_serializable(package='centernet')
 class HourglassBlock(tf.keras.layers.Layer):
   """Hourglass module: an encoder-decoder block."""
 
@@ -274,7 +273,6 @@ class HourglassBlock(tf.keras.layers.Layer):
     return config
 
 
-@tf.keras.utils.register_keras_serializable(package='centernet')
 class CenterNetHeadConv(tf.keras.layers.Layer):
   """Convolution block for the CenterNet head."""
 

official/vision/beta/projects/centernet/modeling/layers/detection_generator.py

@@ -30,7 +30,6 @@ from official.vision.beta.projects.centernet.ops import loss_ops
 from official.vision.beta.projects.centernet.ops import nms_ops
 
 
-@tf.keras.utils.register_keras_serializable(package='centernet')
 class CenterNetDetectionGenerator(tf.keras.layers.Layer):
   """CenterNet Detection Generator."""
 

official/vision/beta/projects/deepmac_maskrcnn/modeling/heads/instance_heads.py

@@ -23,7 +23,6 @@ from official.modeling import tf_utils
 from official.vision.beta.projects.deepmac_maskrcnn.modeling.heads import hourglass_network
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
 class DeepMaskHead(tf.keras.layers.Layer):
   """Creates a mask head."""
 

official/vision/beta/projects/deepmac_maskrcnn/modeling/maskrcnn_model.py

@@ -31,7 +31,6 @@ def resize_as(source, size):
   return tf.transpose(source, (0, 3, 1, 2))
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
 class DeepMaskRCNNModel(maskrcnn_model.MaskRCNNModel):
   """The Mask R-CNN model."""
 

official/vision/beta/projects/example/example_model.py

@@ -25,7 +25,6 @@ import tensorflow as tf
 from official.vision.beta.projects.example import example_config as example_cfg
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
 class ExampleModel(tf.keras.Model):
   """A example model class.
 

official/vision/beta/projects/movinet/modeling/movinet_layers.py

@@ -22,6 +22,7 @@ from typing import Any, Mapping, Optional, Sequence, Tuple, Union
 
 import tensorflow as tf
 
+from official.modeling import tf_utils
 from official.vision.beta.modeling.layers import nn_layers
 
 # Default kernel weight decay that may be overridden
@@ -323,7 +324,8 @@ class ConvBlock(tf.keras.layers.Layer):
     self._use_buffered_input = use_buffered_input
 
     if activation is not None:
-      self._activation_layer = tf.keras.layers.Activation(activation)
+      self._activation_layer = tf_utils.get_activation(
+          activation, use_keras_layer=True)
     else:
       self._activation_layer = None
 

official/vision/beta/projects/movinet/modeling/movinet_layers_test.py

@@ -338,7 +338,7 @@ class MovinetLayersTest(parameterized.TestCase, tf.test.TestCase):
       predicted = tf.concat(predicted, axis=1)
 
       self.assertEqual(predicted.shape, expected.shape)
-      self.assertAllClose(predicted, expected)
+      self.assertAllClose(predicted, expected, atol=1e-4)
 
       self.assertAllClose(
           predicted,
@@ -349,7 +349,8 @@ class MovinetLayersTest(parameterized.TestCase, tf.test.TestCase):
             [[[3., 3., 3.]],
              [[3., 3., 3.]]],
             [[[4., 4., 4.]],
-             [[4., 4., 4.]]]]])
+             [[4., 4., 4.]]]]],
+          atol=1e-4)
 
   def test_stream_movinet_block(self):
     block = movinet_layers.MovinetBlock(

official/vision/beta/projects/panoptic_maskrcnn/README.md

@@ -5,7 +5,7 @@
 Panoptic Segmentation combines the two distinct vision tasks - semantic
 segmentation and instance segmentation. These tasks are unified such that, each
 pixel in the image is assigned the label of the class it belongs to, and also
-the instance identifier of the object it a part of.
+the instance identifier of the object it is a part of.
 
 ## Environment setup
 The code can be run on multiple GPUs or TPUs with different distribution
@@ -13,8 +13,85 @@ strategies. See the TensorFlow distributed training
 [guide](https://www.tensorflow.org/guide/distributed_training) for an overview
 of `tf.distribute`.
 
-The code is compatible with TensorFlow 2.4+. See requirements.txt for all
-prerequisites, and you can also install them using the following command. `pip
-install -r ./official/requirements.txt`
+The code is compatible with TensorFlow 2.6+. See requirements.txt for all
+prerequisites.
 
-**DISCLAIMER**: Panoptic MaskRCNN is still under active development, stay tuned!
+```bash
+$ git clone https://github.com/tensorflow/models.git
+$ cd models
+$ pip3 install -r official/requirements.txt
+$ export PYTHONPATH=$(pwd)
+```
+
+## Preparing Dataset
+```bash
+$ ./official/vision/beta/data/process_coco_panoptic.sh <path-to-data-directory>
+```
+
+## Launch Training
+```bash
+$ export MODEL_DIR="gs://<path-to-model-directory>"
+$ export TPU_NAME="<tpu-name>"
+$ export ANNOTATION_FILE="gs://<path-to-coco-annotation-json>"
+$ export TRAIN_DATA="gs://<path-to-train-data>"
+$ export EVAL_DATA="gs://<path-to-eval-data>"
+$ export OVERRIDES="task.validation_data.input_path=${EVAL_DATA},\
+task.train_data.input_path=${TRAIN_DATA},\
+task.annotation_file=${ANNOTATION_FILE},\
+runtime.distribution_strategy=tpu"
+
+
+$ python3 train.py \
+  --experiment panoptic_fpn_coco \
+  --config_file configs/experiments/r50fpn_1x_coco.yaml \
+  --mode train \
+  --model_dir $MODEL_DIR \
+  --tpu $TPU_NAME \
+  --params_override=$OVERRIDES
+```
+
+## Launch Evaluation
+```bash
+$ export MODEL_DIR="gs://<path-to-model-directory>"
+$ export NUM_GPUS="<number-of-gpus>"
+$ export PRECISION="<floating-point-precision>"
+$ export ANNOTATION_FILE="gs://<path-to-coco-annotation-json>"
+$ export TRAIN_DATA="gs://<path-to-train-data>"
+$ export EVAL_DATA="gs://<path-to-eval-data>"
+$ export OVERRIDES="task.validation_data.input_path=${EVAL_DATA}, \
+task.train_data.input_path=${TRAIN_DATA}, \
+task.annotation_file=${ANNOTATION_FILE}, \
+runtime.distribution_strategy=mirrored, \
+runtime.mixed_precision_dtype=$PRECISION, \
+runtime.num_gpus=$NUM_GPUS"
+
+
+$ python3 train.py \
+  --experiment panoptic_fpn_coco \
+  --config_file configs/experiments/r50fpn_1x_coco.yaml \
+  --mode eval \
+  --model_dir $MODEL_DIR \
+  --params_override=$OVERRIDES
+```
+**Note**: The [PanopticSegmentationGenerator](https://github.com/tensorflow/models/blob/ac7f9e7f2d0508913947242bad3e23ef7cae5a43/official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/panoptic_segmentation_generator.py#L22) layer uses dynamic shapes and hence generating panoptic masks is not supported on Cloud TPUs. Running evaluation on Cloud TPUs is not supported for the same reason. However, training is supported on both Cloud TPUs and GPUs.
+## Pretrained Models
+### Panoptic FPN
+Backbone     | Schedule     | Experiment name             | Box mAP |  Mask mAP  | Overall PQ | Things PQ | Stuff PQ | Checkpoints
+:------------| :----------- | :---------------------------| ------- | ---------- | ---------- | --------- | -------- | ------------:
+ResNet-50    | 1x           | `panoptic_fpn_coco`         | 38.19   |   34.25    |   39.14    |  45.42    |  29.65   | [ckpt](gs://tf_model_garden/vision/panoptic/panoptic_fpn/panoptic_fpn_1x)
+ResNet-50    | 3x           | `panoptic_fpn_coco`         | 40.64   |   36.29    |   40.91    |  47.68    |  30.69   | [ckpt](gs://tf_model_garden/vision/panoptic/panoptic_fpn/panoptic_fpn_3x)
+
+**Note**: Here 1x schedule refers to ~12 epochs
+
+___
+## Citation
+```
+@misc{kirillov2019panoptic,
+      title={Panoptic Feature Pyramid Networks},
+      author={Alexander Kirillov and Ross Girshick and Kaiming He and Piotr Dollár},
+      year={2019},
+      eprint={1901.02446},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

official/vision/beta/projects/panoptic_maskrcnn/configs/experiments/r50fpn_1x_coco.yaml

+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'bfloat16'
+
+task:
+  init_checkpoint: 'gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080'
+  annotation_file: 'coco/instances_val2017.json'
+  train_data:
+    global_batch_size: 64
+  validation_data:
+    global_batch_size: 8
+
+trainer:
+  train_steps: 22500
+  optimizer_config:
+    learning_rate:
+      type: 'stepwise'
+      stepwise:
+        boundaries: [15000, 20000]
+        values: [0.12, 0.012, 0.0012]
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 500
+        warmup_learning_rate: 0.0067

official/vision/beta/projects/panoptic_maskrcnn/configs/experiments/r50fpn_3x_coco.yaml

+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'bfloat16'
+
+task:
+  init_checkpoint: 'gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080'
+  annotation_file: 'coco/instances_val2017.json'
+  train_data:
+    global_batch_size: 64
+  validation_data:
+    global_batch_size: 8
+
+trainer:
+  train_steps: 67500
+  optimizer_config:
+    learning_rate:
+      type: 'stepwise'
+      stepwise:
+        boundaries: [45000, 60000]
+        values: [0.12, 0.012, 0.0012]
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 500
+        warmup_learning_rate: 0.0067

official/vision/beta/projects/panoptic_maskrcnn/modeling/panoptic_maskrcnn_model.py

@@ -21,7 +21,6 @@ import tensorflow as tf
 from official.vision.beta.modeling import maskrcnn_model
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
 class PanopticMaskRCNNModel(maskrcnn_model.MaskRCNNModel):
   """The Panoptic Segmentation model."""
 
@@ -171,7 +170,7 @@ class PanopticMaskRCNNModel(maskrcnn_model.MaskRCNNModel):
       decoder_features = model_outputs['decoder_features']
 
     segmentation_outputs = self.segmentation_head(
-        backbone_features, decoder_features, training=training)
+        (backbone_features, decoder_features), training=training)
 
     model_outputs.update({
         'segmentation_outputs': segmentation_outputs,

official/vision/beta/projects/panoptic_maskrcnn/serving/panoptic_segmentation.py

@@ -18,7 +18,7 @@ from typing import List
 
 import tensorflow as tf
 
-from official.modeling.hyperparams import config_definitions as cfg
+from official.core import config_definitions as cfg
 from official.vision.beta.projects.panoptic_maskrcnn.modeling import panoptic_maskrcnn_model
 from official.vision.beta.serving import detection
 

official/vision/beta/projects/panoptic_maskrcnn/tasks/panoptic_maskrcnn.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 
 """Panoptic MaskRCNN task definition."""
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, List, Mapping, Optional, Tuple, Dict
 from absl import logging
 import tensorflow as tf
 
 from official.common import dataset_fn
 from official.core import task_factory
 from official.vision.beta.dataloaders import input_reader_factory
-from official.vision.beta.evaluation import coco_evaluator
 from official.vision.beta.evaluation import panoptic_quality_evaluator
 from official.vision.beta.evaluation import segmentation_metrics
 from official.vision.beta.losses import segmentation_losses
@@ -235,10 +234,7 @@ class PanopticMaskRCNNTask(maskrcnn.MaskRCNNTask):
             dtype=tf.float32)
 
     else:
-      self.coco_metric = coco_evaluator.COCOEvaluator(
-          annotation_file=self.task_config.annotation_file,
-          include_mask=self.task_config.model.include_mask,
-          per_category_metrics=self.task_config.per_category_metrics)
+      self._build_coco_metrics()
 
       rescale_predictions = (not self.task_config.validation_data.parser
                              .segmentation_resize_eval_groundtruth)
@@ -379,12 +375,12 @@ class PanopticMaskRCNNTask(maskrcnn.MaskRCNNTask):
         'image_info': labels['image_info']
     }
 
+    logs.update(
+        {self.coco_metric.name: (labels['groundtruths'], coco_model_outputs)})
     if self._process_iou_metric_on_cpu:
       logs.update({
-          self.coco_metric.name: (labels['groundtruths'], coco_model_outputs),
-          self.segmentation_perclass_iou_metric.name: (
-              segmentation_labels,
-              outputs['segmentation_outputs'])
+          self.segmentation_perclass_iou_metric.name:
+              (segmentation_labels, outputs['segmentation_outputs'])
       })
     else:
       self.segmentation_perclass_iou_metric.update_state(
@@ -430,7 +426,7 @@ class PanopticMaskRCNNTask(maskrcnn.MaskRCNNTask):
 
   def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
     result = {}
-    result[self.coco_metric.name] = super(
+    result = super(
         PanopticMaskRCNNTask, self).reduce_aggregated_logs(
             aggregated_logs=aggregated_logs,
             global_step=global_step)

official/vision/beta/projects/simclr/heads/simclr_head.py

@@ -24,7 +24,6 @@ regularizers = tf.keras.regularizers
 layers = tf.keras.layers
 
 
-@tf.keras.utils.register_keras_serializable(package='simclr')
 class ProjectionHead(tf.keras.layers.Layer):
   """Projection head."""
 
@@ -144,7 +143,6 @@ class ProjectionHead(tf.keras.layers.Layer):
     return proj_head_output, proj_finetune_output
 
 
-@tf.keras.utils.register_keras_serializable(package='simclr')
 class ClassificationHead(tf.keras.layers.Layer):
   """Classification Head."""
 

official/vision/beta/projects/simclr/losses/contrastive_losses.py

@@ -57,7 +57,6 @@ def cross_replica_concat(tensor: tf.Tensor, num_replicas: int) -> tf.Tensor:
     return tf.reshape(ext_tensor, [-1] + ext_tensor.shape.as_list()[2:])
 
 
-@tf.keras.utils.register_keras_serializable(package='simclr')
 class ContrastiveLoss(object):
   """Contrastive training loss function."""
 

official/vision/beta/projects/simclr/modeling/layers/nn_blocks.py

@@ -22,7 +22,6 @@ from official.modeling import tf_utils
 regularizers = tf.keras.regularizers
 
 
-@tf.keras.utils.register_keras_serializable(package='simclr')
 class DenseBN(tf.keras.layers.Layer):
   """Modified Dense layer to help build simclr system.
 

official/vision/beta/projects/simclr/modeling/simclr_model.py

@@ -27,7 +27,6 @@ PROJECTION_OUTPUT_KEY = 'projection_outputs'
 SUPERVISED_OUTPUT_KEY = 'supervised_outputs'
 
 
-@tf.keras.utils.register_keras_serializable(package='simclr')
 class SimCLRModel(tf.keras.Model):
   """A classification model based on SimCLR framework."""
 

official/vision/beta/projects/video_ssl/modeling/video_ssl_model.py

@@ -27,7 +27,6 @@ from official.vision.beta.projects.video_ssl.configs import video_ssl as video_s
 layers = tf.keras.layers
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
 class VideoSSLModel(tf.keras.Model):
   """A video ssl model class builder."""
 

official/vision/beta/projects/yolo/configs/decoders.py

@@ -33,6 +33,8 @@ class YoloDecoder(hyperparams.Config):
   use_separable_conv: bool = False
   csp_stack: Optional[bool] = None
   fpn_depth: Optional[int] = None
+  max_fpn_depth: Optional[int] = None
+  max_csp_stack: Optional[int] = None
   fpn_filter_scale: Optional[int] = None
   path_process_len: Optional[int] = None
   max_level_process_len: Optional[int] = None

official/vision/beta/projects/yolo/configs/experiments/scaled-yolo/detection/yolo_csp_640_tpu.yaml

+# --experiment_type=scaled_yolo
+# mAP 47.6
+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'float32'
+  tpu_enable_xla_dynamic_padder: false
+task:
+  model:
+    input_size: [640, 640, 3]
+    backbone:
+      type: 'darknet'
+      darknet:
+        model_id: 'altered_cspdarknet53'
+        max_level: 5
+        min_level: 3
+    decoder:
+      type: yolo_decoder
+      yolo_decoder:
+        version: v4
+        type: csp
+    head:
+      smart_bias: true
+    detection_generator:
+      box_type:
+        'all': scaled
+      scale_xy:
+        'all': 2.0
+      max_boxes: 300
+      nms_type: iou
+      iou_thresh: 0.001
+      nms_thresh: 0.65
+    loss:
+      use_scaled_loss: true
+      update_on_repeat: true
+      box_loss_type:
+        'all': ciou
+      ignore_thresh:
+        'all': 0.0
+      iou_normalizer:
+        'all': 0.05
+      cls_normalizer:
+        'all': 0.3
+      object_normalizer:
+        '5': 0.28
+        '4': 0.70
+        '3': 2.80
+      objectness_smooth:
+        'all': 1.0
+    norm_activation:
+      use_sync_bn: true
+    num_classes: 80
+    anchor_boxes:
+      anchors_per_scale: 3
+      boxes: [box: [12, 16], box: [19, 36], box: [40, 28],
+              box: [36, 75], box: [76, 55], box: [72, 146],
+              box: [142, 110], box: [192, 243], box: [459, 401]]
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/tensorflow-performance-data/datasets/coco/train*'
+    shuffle_buffer_size: 10000
+    parser:
+      mosaic:
+        mosaic_frequency: 1.0
+        mixup_frequency: 0.0
+        mosaic_crop_mode: 'scale'
+        mosaic_center: 0.25
+        aug_scale_min: 0.1
+        aug_scale_max: 1.9
+      max_num_instances: 300
+      letter_box: true
+      random_flip: true
+      aug_rand_translate: 0.1
+      area_thresh: 0.1
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/tensorflow-performance-data/datasets/coco/val*'