Internal change

PiperOrigin-RevId: 374481243

official/vision/beta/configs/experiments/maskrcnn/coco_spinenet143_casrcnn_tpu.yaml → official/vision/beta/configs/experiments/maskrcnn/coco_spinenet143_cascadercnn_tpu.yaml

-# Expect to reach: box mAP: 51.6%, mask mAP: 44.5% on COCO
+# --experiment_type=cascadercnn_spinenet_coco
+# Expect to reach: box mAP: 51.9%, mask mAP: 45.0% on COCO
 runtime:
   distribution_strategy: 'tpu'
   mixed_precision_dtype: 'bfloat16'
@@ -43,12 +44,12 @@ task:
     detection_generator:
       pre_nms_top_k: 1000
 trainer:
-  train_steps: 162050
+  train_steps: 231000
   optimizer_config:
     learning_rate:
       type: 'stepwise'
       stepwise:
-        boundaries: [148160, 157420]
+        boundaries: [219450, 226380]
         values: [0.32, 0.032, 0.0032]
     warmup:
       type: 'linear'

official/vision/beta/configs/experiments/maskrcnn/coco_spinenet49_casrcnn_tpu.yaml → official/vision/beta/configs/experiments/maskrcnn/coco_spinenet49_cascadercnn_tpu.yaml

+# --experiment_type=cascadercnn_spinenet_coco
+# Expect to reach: box mAP: 46.4%, mask mAP: 40.0% on COCO
 runtime:
   distribution_strategy: 'tpu'
   mixed_precision_dtype: 'bfloat16'
@@ -48,7 +50,7 @@ trainer:
       type: 'stepwise'
       stepwise:
         boundaries: [219450, 226380]
-        values: [0.28, 0.028, 0.0028]
+        values: [0.32, 0.032, 0.0032]
     warmup:
       type: 'linear'
       linear:

official/vision/beta/configs/maskrcnn.py

@@ -233,6 +233,8 @@ def fasterrcnn_resnetfpn_coco() -> cfg.ExperimentConfig:
   """COCO object detection with Faster R-CNN."""
   steps_per_epoch = 500
   coco_val_samples = 5000
+  train_batch_size = 64
+  eval_batch_size = 8
 
   config = cfg.ExperimentConfig(
       runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
@@ -252,16 +254,16 @@ def fasterrcnn_resnetfpn_coco() -> cfg.ExperimentConfig:
           train_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
               is_training=True,
-              global_batch_size=64,
+              global_batch_size=train_batch_size,
               parser=Parser(
                   aug_rand_hflip=True, aug_scale_min=0.8, aug_scale_max=1.25)),
           validation_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
               is_training=False,
-              global_batch_size=8)),
+              global_batch_size=eval_batch_size)),
       trainer=cfg.TrainerConfig(
           train_steps=22500,
-          validation_steps=coco_val_samples // 8,
+          validation_steps=coco_val_samples // eval_batch_size,
           validation_interval=steps_per_epoch,
           steps_per_loop=steps_per_epoch,
           summary_interval=steps_per_epoch,
@@ -300,6 +302,8 @@ def maskrcnn_resnetfpn_coco() -> cfg.ExperimentConfig:
   """COCO object detection with Mask R-CNN."""
   steps_per_epoch = 500
   coco_val_samples = 5000
+  train_batch_size = 64
+  eval_batch_size = 8
 
   config = cfg.ExperimentConfig(
       runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
@@ -314,16 +318,16 @@ def maskrcnn_resnetfpn_coco() -> cfg.ExperimentConfig:
           train_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
               is_training=True,
-              global_batch_size=64,
+              global_batch_size=train_batch_size,
               parser=Parser(
                   aug_rand_hflip=True, aug_scale_min=0.8, aug_scale_max=1.25)),
           validation_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
               is_training=False,
-              global_batch_size=8)),
+              global_batch_size=eval_batch_size)),
       trainer=cfg.TrainerConfig(
           train_steps=22500,
-          validation_steps=coco_val_samples // 8,
+          validation_steps=coco_val_samples // eval_batch_size,
           validation_interval=steps_per_epoch,
           steps_per_loop=steps_per_epoch,
           summary_interval=steps_per_epoch,
@@ -357,40 +361,50 @@ def maskrcnn_resnetfpn_coco() -> cfg.ExperimentConfig:
   return config
 
 
-@exp_factory.register_config_factory('cascadercnn_resnetfpn_coco')
-def cascadercnn_resnetfpn_coco() -> cfg.ExperimentConfig:
-  """COCO object detection with Cascade R-CNN."""
-  steps_per_epoch = 500
+@exp_factory.register_config_factory('maskrcnn_spinenet_coco')
+def maskrcnn_spinenet_coco() -> cfg.ExperimentConfig:
+  """COCO object detection with Mask R-CNN with SpineNet backbone."""
+  steps_per_epoch = 463
   coco_val_samples = 5000
+  train_batch_size = 256
+  eval_batch_size = 8
 
   config = cfg.ExperimentConfig(
       runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
       task=MaskRCNNTask(
-          init_checkpoint='gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080',
-          init_checkpoint_modules='backbone',
           annotation_file=os.path.join(COCO_INPUT_PATH_BASE,
                                        'instances_val2017.json'),
           model=MaskRCNN(
+              backbone=backbones.Backbone(
+                  type='spinenet',
+                  spinenet=backbones.SpineNet(
+                      model_id='49',
+                      min_level=3,
+                      max_level=7,
+                  )),
+              decoder=decoders.Decoder(
+                  type='identity', identity=decoders.Identity()),
+              anchor=Anchor(anchor_size=3),
+              norm_activation=common.NormActivation(use_sync_bn=True),
               num_classes=91,
-              input_size=[1024, 1024, 3],
-              include_mask=True,
-              roi_sampler=ROISampler(cascade_iou_thresholds=[0.6, 0.7]),
-              detection_head=DetectionHead(
-                  class_agnostic_bbox_pred=True, cascade_class_ensemble=True)),
+              input_size=[640, 640, 3],
+              min_level=3,
+              max_level=7,
+              include_mask=True),
           losses=Losses(l2_weight_decay=0.00004),
           train_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
               is_training=True,
-              global_batch_size=64,
+              global_batch_size=train_batch_size,
               parser=Parser(
-                  aug_rand_hflip=True, aug_scale_min=0.8, aug_scale_max=1.25)),
+                  aug_rand_hflip=True, aug_scale_min=0.5, aug_scale_max=2.0)),
           validation_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
               is_training=False,
-              global_batch_size=8)),
+              global_batch_size=eval_batch_size)),
       trainer=cfg.TrainerConfig(
-          train_steps=22500,
-          validation_steps=coco_val_samples // 8,
+          train_steps=steps_per_epoch * 350,
+          validation_steps=coco_val_samples // eval_batch_size,
           validation_interval=steps_per_epoch,
           steps_per_loop=steps_per_epoch,
           summary_interval=steps_per_epoch,
@@ -405,30 +419,36 @@ def cascadercnn_resnetfpn_coco() -> cfg.ExperimentConfig:
               'learning_rate': {
                   'type': 'stepwise',
                   'stepwise': {
-                      'boundaries': [15000, 20000],
-                      'values': [0.12, 0.012, 0.0012],
+                      'boundaries': [
+                          steps_per_epoch * 320, steps_per_epoch * 340
+                      ],
+                      'values': [0.32, 0.032, 0.0032],
                   }
               },
               'warmup': {
                   'type': 'linear',
                   'linear': {
-                      'warmup_steps': 500,
+                      'warmup_steps': 2000,
                       'warmup_learning_rate': 0.0067
                   }
               }
           })),
       restrictions=[
           'task.train_data.is_training != None',
-          'task.validation_data.is_training != None'
+          'task.validation_data.is_training != None',
+          'task.model.min_level == task.model.backbone.spinenet.min_level',
+          'task.model.max_level == task.model.backbone.spinenet.max_level',
       ])
   return config
 
 
-@exp_factory.register_config_factory('maskrcnn_spinenet_coco')
-def maskrcnn_spinenet_coco() -> cfg.ExperimentConfig:
-  """COCO object detection with Mask R-CNN with SpineNet backbone."""
+@exp_factory.register_config_factory('cascadercnn_spinenet_coco')
+def cascadercnn_spinenet_coco() -> cfg.ExperimentConfig:
+  """COCO object detection with Cascade R-CNN with SpineNet backbone."""
   steps_per_epoch = 463
   coco_val_samples = 5000
+  train_batch_size = 256
+  eval_batch_size = 8
 
   config = cfg.ExperimentConfig(
       runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
@@ -445,8 +465,12 @@ def maskrcnn_spinenet_coco() -> cfg.ExperimentConfig:
                   )),
               decoder=decoders.Decoder(
                   type='identity', identity=decoders.Identity()),
+              roi_sampler=ROISampler(cascade_iou_thresholds=[0.6, 0.7]),
+              detection_head=DetectionHead(
+                  class_agnostic_bbox_pred=True, cascade_class_ensemble=True),
               anchor=Anchor(anchor_size=3),
-              norm_activation=common.NormActivation(use_sync_bn=True),
+              norm_activation=common.NormActivation(
+                  use_sync_bn=True, activation='swish'),
               num_classes=91,
               input_size=[640, 640, 3],
               min_level=3,
@@ -456,16 +480,16 @@ def maskrcnn_spinenet_coco() -> cfg.ExperimentConfig:
           train_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
               is_training=True,
-              global_batch_size=256,
+              global_batch_size=train_batch_size,
               parser=Parser(
-                  aug_rand_hflip=True, aug_scale_min=0.5, aug_scale_max=2.0)),
+                  aug_rand_hflip=True, aug_scale_min=0.1, aug_scale_max=2.5)),
           validation_data=DataConfig(
               input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
               is_training=False,
-              global_batch_size=8)),
+              global_batch_size=eval_batch_size)),
       trainer=cfg.TrainerConfig(
-          train_steps=steps_per_epoch * 350,
-          validation_steps=coco_val_samples // 8,
+          train_steps=steps_per_epoch * 500,
+          validation_steps=coco_val_samples // eval_batch_size,
           validation_interval=steps_per_epoch,
           steps_per_loop=steps_per_epoch,
           summary_interval=steps_per_epoch,
@@ -481,9 +505,9 @@ def maskrcnn_spinenet_coco() -> cfg.ExperimentConfig:
                   'type': 'stepwise',
                   'stepwise': {
                       'boundaries': [
-                          steps_per_epoch * 320, steps_per_epoch * 340
+                          steps_per_epoch * 475, steps_per_epoch * 490
                       ],
-                      'values': [0.28, 0.028, 0.0028],
+                      'values': [0.32, 0.032, 0.0032],
                   }
               },
               'warmup': {

official/vision/beta/configs/maskrcnn_test.py

@@ -29,7 +29,7 @@ class MaskRCNNConfigTest(tf.test.TestCase, parameterized.TestCase):
       ('fasterrcnn_resnetfpn_coco',),
       ('maskrcnn_resnetfpn_coco',),
       ('maskrcnn_spinenet_coco',),
-      ('cascadercnn_resnetfpn_coco',),
+      ('cascadercnn_spinenet_coco',),
   )
   def test_maskrcnn_configs(self, config_name):
     config = exp_factory.get_exp_config(config_name)