Internal change

PiperOrigin-RevId: 520733955

Author: Zarjagen

official/projects/deepmac_maskrcnn/tasks/deep_mask_head_rcnn.py

@@ -163,6 +163,13 @@ def build_maskrcnn(input_specs: tf.keras.layers.InputSpec,
       mask_head=mask_head,
       mask_sampler=mask_sampler_obj,
       mask_roi_aligner=mask_roi_aligner_obj,
+      class_agnostic_bbox_pred=detection_head_config.class_agnostic_bbox_pred,
+      cascade_class_ensemble=detection_head_config.cascade_class_ensemble,
+      min_level=model_config.min_level,
+      max_level=model_config.max_level,
+      num_scales=model_config.anchor.num_scales,
+      aspect_ratios=model_config.anchor.aspect_ratios,
+      anchor_size=model_config.anchor.anchor_size,
       outer_boxes_scale=model_config.outer_boxes_scale,
       use_gt_boxes_for_masks=model_config.use_gt_boxes_for_masks)
   return model
@@ -193,4 +200,9 @@ def build_model(self):
     if self.task_config.freeze_backbone:
       model.backbone.trainable = False
 
+    # Builds the model through warm-up call.
+    dummy_images = tf.keras.Input(self.task_config.model.input_size)
+    dummy_image_shape = tf.keras.layers.Input([2])
+    _ = model(dummy_images, image_shape=dummy_image_shape, training=False)
+
     return model

official/projects/panoptic/configs/panoptic_deeplab.py

@@ -129,7 +129,8 @@ class PanopticDeeplab(hyperparams.Config):
   norm_activation: common.NormActivation = common.NormActivation()
   backbone: backbones.Backbone = backbones.Backbone(
       type='resnet', resnet=backbones.ResNet())
-  decoder: decoders.Decoder = decoders.Decoder(type='aspp')
+  decoder: decoders.Decoder = decoders.Decoder(
+      type='aspp', aspp=decoders.ASPP(level=3))
   semantic_head: SemanticHead = SemanticHead()
   instance_head: InstanceHead = InstanceHead()
   shared_decoder: bool = False

official/projects/panoptic/configs/panoptic_maskrcnn.py

@@ -38,6 +38,7 @@
 _COCO_VAL_EXAMPLES = 5000
 
 # pytype: disable=wrong-keyword-args
+# pylint: disable=unexpected-keyword-arg
 
 
 @dataclasses.dataclass
@@ -108,10 +109,9 @@ class Backbone(backbones.Backbone):
 @dataclasses.dataclass
 class PanopticMaskRCNN(deepmac_maskrcnn.DeepMaskHeadRCNN):
   """Panoptic Mask R-CNN model config."""
-  backbone: Backbone = Backbone()
-  segmentation_model: semantic_segmentation.SemanticSegmentationModel = (
-      SEGMENTATION_MODEL(num_classes=2))
-  include_mask = True
+  backbone: Backbone = Backbone(type='resnet', resnet=backbones.ResNet())
+  segmentation_model: SEGMENTATION_MODEL = SEGMENTATION_MODEL(num_classes=2)
+  include_mask: bool = True
   shared_backbone: bool = True
   shared_decoder: bool = True
   stuff_classes_offset: int = 0

official/projects/panoptic/tasks/panoptic_deeplab.py

@@ -50,6 +50,12 @@ def build_model(self):
         input_specs=input_specs,
         model_config=self.task_config.model,
         l2_regularizer=l2_regularizer)
+
+    # Builds the model through warm-up call.
+    dummy_images = tf.keras.Input(self.task_config.model.input_size)
+    # Note that image_info is always in the shape of [4, 2].
+    dummy_image_info = tf.keras.layers.Input([4, 2])
+    _ = model(dummy_images, dummy_image_info, training=False)
     return model
 
   def initialize(self, model: tf.keras.Model):

official/projects/panoptic/tasks/panoptic_maskrcnn.py

@@ -71,6 +71,12 @@ def build_model(self) -> tf.keras.Model:
     if self.task_config.freeze_backbone:
       model.backbone.trainable = False
 
+    # Builds the model through warm-up call.
+    dummy_images = tf.keras.Input(self.task_config.model.input_size)
+    # Note that image_info is always in the shape of [4, 2].
+    dummy_image_info = tf.keras.layers.Input([4, 2])
+    _ = model(dummy_images, image_info=dummy_image_info, training=False)
+
     return model
 
   def initialize(self, model: tf.keras.Model) -> None:

official/vision/tasks/maskrcnn.py

@@ -91,6 +91,11 @@ def build_model(self):
     if self.task_config.freeze_backbone:
       model.backbone.trainable = False
 
+    # Builds the model through warm-up call.
+    dummy_images = tf.keras.Input(self.task_config.model.input_size)
+    dummy_image_shape = tf.keras.layers.Input([2])
+    _ = model(dummy_images, image_shape=dummy_image_shape, training=False)
+
     return model
 
   def initialize(self, model: tf.keras.Model):
@@ -487,7 +492,6 @@ def validation_step(self,
       A dictionary of logs.
     """
     images, labels = inputs
-
     outputs = model(
         images,
         anchor_boxes=labels['anchor_boxes'],