tf-models-nightly 2.18.0.dev20240820__py2.py3-none-any.whl → 2.18.0.dev20240822__py2.py3-none-any.whl

official/projects/maskconver/tasks/maskconver.py

@@ -0,0 +1,641 @@
+# Copyright 2024 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.
+
+"""Panoptic MaskRCNN task definition."""
+from typing import Any, Dict, List, Mapping, Optional, Tuple
+
+import tensorflow as tf, tf_keras
+
+from official.common import dataset_fn
+from official.core import task_factory
+from official.projects.centernet.ops import loss_ops
+from official.projects.maskconver.configs import maskconver as exp_cfg
+from official.projects.maskconver.dataloaders import maskconver_segmentation_input
+from official.projects.maskconver.dataloaders import panoptic_maskrcnn_input
+from official.projects.maskconver.losses import maskconver_losses
+from official.projects.maskconver.modeling import factory
+from official.projects.volumetric_models.losses import segmentation_losses as volumeteric_segmentation_losses
+from official.vision.dataloaders import input_reader_factory
+from official.vision.dataloaders import segmentation_input
+from official.vision.evaluation import panoptic_quality_evaluator
+from official.vision.evaluation import segmentation_metrics
+from official.vision.tasks import maskrcnn
+from official.vision.tasks import semantic_segmentation
+
+
+@task_factory.register_task_cls(exp_cfg.MaskConverTask)
+class PanopticMaskRCNNTask(maskrcnn.MaskRCNNTask):
+
+  """A single-replica view of training procedure.
+
+  Panoptic Mask R-CNN task provides artifacts for training/evalution procedures,
+  including loading/iterating over Datasets, initializing the model, calculating
+  the loss, post-processing, and customized metrics with reduction.
+  """
+
+  def build_model(self) -> tf_keras.Model:
+    """Build Panoptic Mask R-CNN model."""
+
+    tf_keras.utils.set_random_seed(0)
+    tf.config.experimental.enable_op_determinism()
+    input_specs = tf_keras.layers.InputSpec(
+        shape=[None] + self.task_config.model.input_size)
+
+    l2_weight_decay = self.task_config.losses.l2_weight_decay
+    # Divide weight decay by 2.0 to match the implementation of tf.nn.l2_loss.
+    # (https://www.tensorflow.org/api_docs/python/tf/keras/regularizers/l2)
+    # (https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss)
+    l2_regularizer = (tf_keras.regularizers.l2(
+        l2_weight_decay / 2.0) if l2_weight_decay else None)
+
+    model = factory.build_maskconver_model(
+        input_specs=input_specs,
+        model_config=self.task_config.model,
+        l2_regularizer=l2_regularizer)
+    return model
+
+  def build_inputs(
+      self,
+      params: exp_cfg.DataConfig,
+      input_context: Optional[tf.distribute.InputContext] = None
+  ) -> tf.data.Dataset:
+    """Build input dataset."""
+    decoder_cfg = params.decoder.get()
+    if params.decoder.type == 'simple_decoder':
+      decoder = panoptic_maskrcnn_input.TfExampleDecoder(
+          regenerate_source_id=decoder_cfg.regenerate_source_id,
+          mask_binarize_threshold=decoder_cfg.mask_binarize_threshold,
+          include_panoptic_masks=decoder_cfg.include_panoptic_masks,
+          panoptic_category_mask_key=decoder_cfg.panoptic_category_mask_key,
+          panoptic_instance_mask_key=decoder_cfg.panoptic_instance_mask_key)
+    else:
+      raise ValueError('Unknown decoder type: {}!'.format(params.decoder.type))
+
+    parser = panoptic_maskrcnn_input.Parser(
+        output_size=self.task_config.model.input_size[:2],
+        min_level=self.task_config.model.min_level,
+        max_level=self.task_config.model.max_level,
+        num_scales=self.task_config.model.anchor.num_scales,
+        aspect_ratios=self.task_config.model.anchor.aspect_ratios,
+        anchor_size=self.task_config.model.anchor.anchor_size,
+        dtype=params.dtype,
+        rpn_match_threshold=params.parser.rpn_match_threshold,
+        rpn_unmatched_threshold=params.parser.rpn_unmatched_threshold,
+        rpn_batch_size_per_im=params.parser.rpn_batch_size_per_im,
+        rpn_fg_fraction=params.parser.rpn_fg_fraction,
+        aug_rand_hflip=params.parser.aug_rand_hflip,
+        aug_scale_min=params.parser.aug_scale_min,
+        aug_scale_max=params.parser.aug_scale_max,
+        skip_crowd_during_training=params.parser.skip_crowd_during_training,
+        max_num_instances=self.task_config.model.num_instances,
+        mask_crop_size=params.parser.mask_crop_size,
+        segmentation_resize_eval_groundtruth=params.parser
+        .segmentation_resize_eval_groundtruth,
+        segmentation_groundtruth_padded_size=params.parser
+        .segmentation_groundtruth_padded_size,
+        segmentation_ignore_label=params.parser.segmentation_ignore_label,
+        panoptic_ignore_label=params.parser.panoptic_ignore_label,
+        include_panoptic_masks=params.parser.include_panoptic_masks,
+        num_panoptic_categories=self.task_config.model.num_classes,
+        num_thing_categories=self.task_config.model.num_thing_classes,
+        level=self.task_config.model.level,
+        gaussian_iou=params.parser.gaussaian_iou,
+        aug_type=params.parser.aug_type,
+        max_num_stuff_centers=params.parser.max_num_stuff_centers)
+
+    reader = input_reader_factory.input_reader_generator(
+        params,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
+        decoder_fn=decoder.decode,
+        parser_fn=parser.parse_fn(params.is_training))
+    dataset = reader.read(input_context=input_context)
+
+    return dataset
+
+  def build_losses(self,
+                   outputs: Mapping[str, Any],
+                   labels: Mapping[str, Any],
+                   aux_losses: Optional[Any] = None,
+                   step=None) -> Dict[str, tf.Tensor]:
+    """Build Panoptic Mask R-CNN losses."""
+    loss_params = self._task_config.losses
+
+    batch_size = tf.cast(tf.shape(labels['num_instances'])[0], tf.float32)
+    center_loss_fn = maskconver_losses.PenaltyReducedLogisticFocalLoss(
+        alpha=loss_params.alpha, beta=loss_params.beta)
+    mask_loss_fn = maskconver_losses.PenaltyReducedLogisticFocalLoss()
+
+    # Calculate center heatmap loss
+    # TODO(arashwan): add valid weights.
+    # output_unpad_image_shapes = labels['image_info'][:, 0, :]
+    # valid_anchor_weights = loss_ops.get_valid_anchor_weights_in_flattened_image(  # pylint: disable=line-too-long
+    #     output_unpad_image_shapes, h, w)
+    # valid_anchor_weights = tf.expand_dims(valid_anchor_weights, 2)
+
+    true_flattened_ct_heatmap = loss_ops.flatten_spatial_dimensions(
+        labels['panoptic_heatmaps'])
+    true_flattened_ct_heatmap = tf.cast(true_flattened_ct_heatmap, tf.float32)
+
+    pred_flattened_ct_heatmap = loss_ops.flatten_spatial_dimensions(
+        outputs['class_heatmaps'])
+    pred_flattened_ct_heatmap = tf.cast(pred_flattened_ct_heatmap, tf.float32)
+    center_padding_mask = 1 - labels['panoptic_padding_mask'][:, :, :, None]
+    center_padding_mask = tf.image.resize(
+        center_padding_mask, tf.shape(
+            labels['panoptic_heatmaps'])[1:3], method='nearest')
+    center_padding_mask = tf.maximum(center_padding_mask, 0.0)
+    center_padding_mask = center_padding_mask * tf.ones_like(labels['panoptic_heatmaps'])
+    weights_flattened_mask = loss_ops.flatten_spatial_dimensions(
+        center_padding_mask)
+    center_loss = center_loss_fn(
+        target_tensor=true_flattened_ct_heatmap,
+        prediction_tensor=pred_flattened_ct_heatmap,
+        weights=weights_flattened_mask)
+
+    center_loss = tf.reduce_sum(
+        center_loss / (labels['num_instances'][:, None, None] + 1.0))  / batch_size
+
+    gt_masks = labels['panoptic_masks']
+    gt_mask_weights = labels['panoptic_mask_weights'][:, None, None, :] * tf.ones_like(gt_masks)
+    panoptic_padding_mask = labels['panoptic_padding_mask'][:, :, :, None] * tf.ones_like(gt_masks)
+
+    true_flattened_masks = loss_ops.flatten_spatial_dimensions(
+        gt_masks)
+    true_flattened_ct_heatmap = tf.cast(true_flattened_ct_heatmap, tf.float32)
+    predicted_masks = tf.cast(outputs['mask_proposal_logits'], tf.float32)
+    predicted_masks = tf.image.resize(
+        predicted_masks, tf.shape(gt_masks)[1:3], method='bilinear')
+    pred_flattened_masks = loss_ops.flatten_spatial_dimensions(predicted_masks)
+    mask_loss = tf.cast(0.0, tf.float32)
+    mask_loss_fn = tf_keras.losses.BinaryCrossentropy(
+        from_logits=True,
+        label_smoothing=0.0,
+        axis=-1,
+        reduction=tf_keras.losses.Reduction.NONE,
+        name='binary_crossentropy')
+    mask_weights = tf.reshape(
+        tf.cast(true_flattened_masks >= 0, tf.float32),
+        [-1, 1]) * tf.reshape(gt_mask_weights, [-1, 1])  * tf.reshape(
+            (1 - panoptic_padding_mask), [-1, 1])
+    mask_loss = mask_loss_fn(
+        tf.reshape(gt_masks, [-1, 1]),
+        tf.reshape(pred_flattened_masks, [-1, 1]),
+        sample_weight=mask_weights)
+    mask_loss = tf.reduce_sum(mask_loss) / (tf.reduce_sum(mask_weights) + 1.0)
+
+    # Dice loss
+    _, h, w, _ = gt_masks.get_shape().as_list()
+    masked_predictions = tf.sigmoid(predicted_masks) * gt_mask_weights * (1 - panoptic_padding_mask)
+    masked_gt_masks = gt_masks * gt_mask_weights * (1 - panoptic_padding_mask)
+
+    masked_predictions = tf.transpose(masked_predictions, [0, 3, 1, 2])
+    masked_predictions = tf.reshape(masked_predictions, [-1, h, w, 1])
+    masked_gt_masks = tf.transpose(masked_gt_masks, [0, 3, 1, 2])
+    masked_gt_masks = tf.reshape(masked_gt_masks, [-1, h, w, 1])
+
+    dice_loss_fn = volumeteric_segmentation_losses.SegmentationLossDiceScore(
+        metric_type='adaptive', axis=(2, 3))
+    dice_loss = dice_loss_fn(logits=masked_predictions, labels=masked_gt_masks)
+
+    total_loss = (center_loss + loss_params.mask_weight * mask_loss + loss_params.mask_weight * dice_loss)
+
+    if aux_losses:
+      total_loss += tf.add_n(aux_losses)
+
+    total_loss = loss_params.loss_weight * total_loss
+
+    losses = {'total_loss': total_loss,
+              'mask_loss': mask_loss,
+              'center_loss': center_loss,
+              'dice_loss': dice_loss}
+    return losses
+
+  def build_metrics(self, training: bool = True) -> List[
+      tf_keras.metrics.Metric]:
+    """Build detection metrics."""
+    metrics = []
+    if training:
+      metric_names = [
+          'total_loss',
+          'center_loss',
+          'mask_loss',
+          'dice_loss',
+      ]
+      for name in metric_names:
+        metrics.append(tf_keras.metrics.Mean(name, dtype=tf.float32))
+
+    else:
+      pq_config = self.task_config.panoptic_quality_evaluator
+      self.panoptic_quality_metric = (
+          panoptic_quality_evaluator.PanopticQualityEvaluator(
+              num_categories=pq_config.num_categories,
+              ignored_label=pq_config.ignored_label,
+              max_instances_per_category=pq_config.max_instances_per_category,
+              offset=pq_config.offset,
+              is_thing=pq_config.is_thing,
+              rescale_predictions=pq_config.rescale_predictions))
+    return metrics
+
+  def train_step(self,
+                 inputs: Tuple[Any, Any],
+                 model: tf_keras.Model,
+                 optimizer: tf_keras.optimizers.Optimizer,
+                 metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
+    """Does forward and backward.
+
+    Args:
+      inputs: a dictionary of input tensors.
+      model: the model, forward pass definition.
+      optimizer: the optimizer for this training step.
+      metrics: a nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    images, labels = inputs
+    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
+
+    with tf.GradientTape() as tape:
+      outputs = model(
+          images,
+          box_indices=labels['panoptic_box_indices'],
+          classes=labels['panoptic_classes'],
+          training=True)
+      outputs = tf.nest.map_structure(
+          lambda x: tf.cast(x, tf.float32), outputs)
+
+      # Computes per-replica loss.
+      losses = self.build_losses(
+          outputs=outputs,
+          labels=labels,
+          aux_losses=model.losses,
+          step=optimizer.iterations)
+      scaled_loss = losses['total_loss'] / num_replicas
+
+      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
+      # scaled for numerical stability.
+      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
+        scaled_loss = optimizer.get_scaled_loss(scaled_loss)
+
+    tvars = model.trainable_variables
+    grads = tape.gradient(scaled_loss, tvars)
+    # Scales back gradient when LossScaleOptimizer is used.
+    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
+      grads = optimizer.get_unscaled_gradients(grads)
+    optimizer.apply_gradients(list(zip(grads, tvars)))
+
+    logs = {self.loss: losses['total_loss']}
+
+    if metrics:
+      for m in metrics:
+        m.update_state(losses[m.name])
+
+    return logs
+
+  def validation_step(self,
+                      inputs: Tuple[Any, Any],
+                      model: tf_keras.Model,
+                      metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
+    """Validatation step.
+
+    Args:
+      inputs: a dictionary of input tensors.
+      model: the keras.Model.
+      metrics: a nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    images, labels = inputs
+
+    outputs = model(
+        images,
+        image_info=labels['image_info'],
+        training=False)
+
+    logs = {self.loss: 0}
+
+    pq_metric_labels = {
+        'category_mask':
+            labels['groundtruths']['gt_panoptic_category_mask'],
+        'instance_mask':
+            labels['groundtruths']['gt_panoptic_instance_mask'],
+        'image_info': labels['image_info']
+    }
+    logs.update({
+        self.panoptic_quality_metric.name:
+            (pq_metric_labels, outputs['panoptic_outputs'])})
+    return logs
+
+  def aggregate_logs(self, state=None, step_outputs=None):
+    if state is None:
+      self.panoptic_quality_metric.reset_states()
+      state = [self.panoptic_quality_metric]
+
+    self.panoptic_quality_metric.update_state(
+        step_outputs[self.panoptic_quality_metric.name][0],
+        step_outputs[self.panoptic_quality_metric.name][1])
+
+    return state
+
+  def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
+    result = {}
+
+    report_per_class_metrics = (
+        self.task_config.panoptic_quality_evaluator.report_per_class_metrics)
+    panoptic_quality_results = self.panoptic_quality_metric.result()
+    for k, value in panoptic_quality_results.items():
+      if k.endswith('per_class'):
+        if report_per_class_metrics:
+          for i, per_class_value in enumerate(value):
+            metric_key = 'panoptic_quality/{}/class_{}'.format(k, i)
+            result[metric_key] = per_class_value
+        else:
+          continue
+      else:
+        result['panoptic_quality/{}'.format(k)] = value
+    return result
+
+
+@task_factory.register_task_cls(exp_cfg.MaskConverSegTask)
+class MaskConverSegmentation(semantic_segmentation.SemanticSegmentationTask):
+
+  """A single-replica view of training procedure.
+
+  MaskConver task provides artifacts for training/evalution procedures,
+  including loading/iterating over Datasets, initializing the model, calculating
+  the loss, post-processing, and customized metrics with reduction.
+  """
+
+  def build_model(self) -> tf_keras.Model:
+    """Build maskconver model."""
+
+    tf_keras.utils.set_random_seed(0)
+    tf.config.experimental.enable_op_determinism()
+    input_specs = tf_keras.layers.InputSpec(
+        shape=[None] + self.task_config.model.input_size)
+
+    l2_weight_decay = self.task_config.losses.l2_weight_decay
+    # Divide weight decay by 2.0 to match the implementation of tf.nn.l2_loss.
+    # (https://www.tensorflow.org/api_docs/python/tf/keras/regularizers/l2)
+    # (https://www.tensorflow.org/api_docs/python/tf/nn/l2_loss)
+    l2_regularizer = (tf_keras.regularizers.l2(
+        l2_weight_decay / 2.0) if l2_weight_decay else None)
+
+    model = factory.build_maskconver_model(
+        input_specs=input_specs,
+        model_config=self.task_config.model,
+        l2_regularizer=l2_regularizer,
+        segmentation_inference=True)
+    return model
+
+  def build_inputs(self,
+                   params: exp_cfg.DataConfig,
+                   input_context: Optional[tf.distribute.InputContext] = None):
+    """Builds classification input."""
+
+    ignore_label = self.task_config.losses.ignore_label
+
+    decoder = segmentation_input.Decoder()
+
+    parser = maskconver_segmentation_input.Parser(
+        output_size=params.output_size,
+        num_classes=self.task_config.model.num_classes,
+        crop_size=params.crop_size,
+        ignore_label=ignore_label,
+        resize_eval_groundtruth=params.resize_eval_groundtruth,
+        groundtruth_padded_size=params.groundtruth_padded_size,
+        aug_scale_min=params.aug_scale_min,
+        aug_scale_max=params.aug_scale_max,
+        aug_rand_hflip=params.aug_rand_hflip,
+        preserve_aspect_ratio=params.preserve_aspect_ratio,
+        level=self.task_config.model.level,
+        aug_type=params.aug_type,
+        max_num_stuff_centers=params.max_num_stuff_centers,
+        dtype=params.dtype)
+
+    reader = input_reader_factory.input_reader_generator(
+        params,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
+        decoder_fn=decoder.decode,
+        parser_fn=parser.parse_fn(params.is_training))
+
+    dataset = reader.read(input_context=input_context)
+
+    return dataset
+
+  def build_losses(self,
+                   outputs: Mapping[str, Any],
+                   labels: Mapping[str, Any],
+                   aux_losses: Optional[Any] = None,
+                   step=None) -> Dict[str, tf.Tensor]:
+    """Build Panoptic Mask R-CNN losses."""
+    loss_params = self._task_config.losses
+
+    # b, h, w, c = outputs['class_heatmaps'].get_shape().as_list()
+    batch_size = tf.cast(tf.shape(labels['num_instances'])[0], tf.float32)
+    center_loss_fn = maskconver_losses.PenaltyReducedLogisticFocalLoss(
+        alpha=loss_params.alpha, beta=loss_params.beta)
+    mask_loss_fn = maskconver_losses.PenaltyReducedLogisticFocalLoss()
+
+    true_flattened_ct_heatmap = loss_ops.flatten_spatial_dimensions(
+        labels['seg_ct_heatmaps'])
+    true_flattened_ct_heatmap = tf.cast(true_flattened_ct_heatmap, tf.float32)
+
+    pred_flattened_ct_heatmap = loss_ops.flatten_spatial_dimensions(
+        outputs['class_heatmaps'])
+    pred_flattened_ct_heatmap = tf.cast(pred_flattened_ct_heatmap, tf.float32)
+    center_valid_mask = labels['seg_valid_mask'][:, :, :, None]
+    center_valid_mask = tf.image.resize(
+        center_valid_mask, tf.shape(
+            labels['seg_ct_heatmaps'])[1:3], method='nearest')
+    center_valid_mask = tf.maximum(center_valid_mask, 0.0)
+    center_valid_mask = center_valid_mask * tf.ones_like(
+        labels['seg_ct_heatmaps'])
+    weights_flattened_mask = loss_ops.flatten_spatial_dimensions(
+        center_valid_mask)
+    center_loss = center_loss_fn(
+        target_tensor=true_flattened_ct_heatmap,
+        prediction_tensor=pred_flattened_ct_heatmap,
+        weights=weights_flattened_mask)
+
+    center_loss = tf.reduce_sum(
+        center_loss /
+        (labels['num_instances'][:, None, None] + 1.0)) / batch_size
+
+    gt_masks = labels['seg_masks']
+    gt_mask_weights = labels['seg_mask_weights'][:, None,
+                                                 None, :] * tf.ones_like(
+                                                     gt_masks)
+    valid_mask = labels['seg_valid_mask'][:, :, :,
+                                          None] * tf.ones_like(gt_masks)
+
+    true_flattened_masks = loss_ops.flatten_spatial_dimensions(gt_masks)
+    true_flattened_ct_heatmap = tf.cast(true_flattened_ct_heatmap, tf.float32)
+    predicted_masks = tf.cast(outputs['mask_proposal_logits'], tf.float32)
+    predicted_masks = tf.image.resize(
+        predicted_masks, tf.shape(gt_masks)[1:3], method='bilinear')
+    pred_flattened_masks = loss_ops.flatten_spatial_dimensions(predicted_masks)
+    mask_loss = tf.cast(0.0, tf.float32)
+
+    mask_loss_fn = tf_keras.losses.BinaryCrossentropy(
+        from_logits=True,
+        label_smoothing=0.0,
+        axis=-1,
+        reduction=tf_keras.losses.Reduction.NONE,
+        name='binary_crossentropy')
+    mask_weights = tf.reshape(
+        tf.cast(true_flattened_masks >= 0, tf.float32),
+        [-1, 1]) * tf.reshape(gt_mask_weights, [-1, 1])  * tf.reshape(
+            (valid_mask), [-1, 1])
+    mask_loss = mask_loss_fn(
+        tf.reshape(gt_masks, [-1, 1]),
+        tf.reshape(pred_flattened_masks, [-1, 1]),
+        sample_weight=mask_weights)
+
+    mask_loss = tf.reduce_sum(mask_loss) / (tf.reduce_sum(mask_weights) + 1.0)
+
+    total_loss = (center_loss + loss_params.mask_weight * mask_loss)
+
+    if aux_losses:
+      total_loss += tf.add_n(aux_losses)
+
+    total_loss = loss_params.loss_weight * total_loss
+
+    losses = {'total_loss': total_loss,
+              'mask_loss': mask_loss,
+              'center_loss': center_loss}
+    return losses
+
+  def build_metrics(self, training: bool = True) -> List[
+      tf_keras.metrics.Metric]:
+    """Build detection metrics."""
+    metrics = []
+    if training:
+      metric_names = [
+          'total_loss',
+          'center_loss',
+          'mask_loss',
+      ]
+      for name in metric_names:
+        metrics.append(tf_keras.metrics.Mean(name, dtype=tf.float32))
+    else:
+      self.iou_metric = segmentation_metrics.PerClassIoU(
+          name='per_class_iou',
+          num_classes=self.task_config.model.num_classes,
+          rescale_predictions=False,
+          dtype=tf.float32)
+
+    return metrics
+
+  def train_step(self,
+                 inputs: Tuple[Any, Any],
+                 model: tf_keras.Model,
+                 optimizer: tf_keras.optimizers.Optimizer,
+                 metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
+    """Does forward and backward.
+
+    Args:
+      inputs: a dictionary of input tensors.
+      model: the model, forward pass definition.
+      optimizer: the optimizer for this training step.
+      metrics: a nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    images, labels = inputs
+    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
+
+    with tf.GradientTape() as tape:
+      outputs = model(
+          images,
+          box_indices=labels['seg_box_indices'],
+          classes=labels['seg_classes'],
+          training=True)
+      outputs = tf.nest.map_structure(
+          lambda x: tf.cast(x, tf.float32), outputs)
+
+      # Computes per-replica loss.
+      losses = self.build_losses(
+          outputs=outputs,
+          labels=labels,
+          aux_losses=model.losses,
+          step=optimizer.iterations)
+      scaled_loss = losses['total_loss'] / num_replicas
+
+      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
+      # scaled for numerical stability.
+      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
+        scaled_loss = optimizer.get_scaled_loss(scaled_loss)
+
+    tvars = model.trainable_variables
+    grads = tape.gradient(scaled_loss, tvars)
+    # Scales back gradient when LossScaleOptimizer is used.
+    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
+      grads = optimizer.get_unscaled_gradients(grads)
+    optimizer.apply_gradients(list(zip(grads, tvars)))
+
+    logs = {self.loss: losses['total_loss']}
+
+    if metrics:
+      for m in metrics:
+        m.update_state(losses[m.name])
+
+    return logs
+
+  def validation_step(self,
+                      inputs: Tuple[Any, Any],
+                      model: tf_keras.Model,
+                      metrics: Optional[List[Any]] = None):
+    """Validatation step.
+
+    Args:
+      inputs: a dictionary of input tensors.
+      model: the keras.Model.
+      metrics: a nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    features, labels = inputs
+
+    outputs = model(
+        features,
+        image_info=labels['image_info'],
+        training=False)
+    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
+
+    logs = {self.loss: 0}
+    outputs = tf.one_hot(
+        tf.cast(outputs['panoptic_outputs']['category_mask'], tf.int32),
+        self.task_config.model.num_classes)
+
+    self.iou_metric.update_state(labels, tf.cast(outputs, tf.float32))
+    return logs
+
+  def aggregate_logs(self, state=None, step_outputs=None):
+    if state is None:
+      self.iou_metric.reset_states()
+      state = self.iou_metric
+    return state
+
+  def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
+    result = {}
+    ious = self.iou_metric.result()
+    for i, value in enumerate(ious.numpy()):
+      result.update({'iou/{}'.format(i): value})
+    # Computes mean IoU
+    result.update({'mean_iou': tf.reduce_mean(ious).numpy()})
+    return result