# Copyright 2022 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 Deeplab task definition.""" from typing import Any, Dict, List, Mapping, Optional, Tuple from absl import logging import tensorflow as tf from official.common import dataset_fn from official.core import base_task from official.core import task_factory from official.projects.panoptic.configs import panoptic_deeplab as exp_cfg from official.projects.panoptic.dataloaders import panoptic_deeplab_input from official.projects.panoptic.losses import panoptic_deeplab_losses from official.projects.panoptic.modeling import factory from official.vision.dataloaders import input_reader_factory from official.vision.evaluation import panoptic_quality_evaluator from official.vision.evaluation import segmentation_metrics @task_factory.register_task_cls(exp_cfg.PanopticDeeplabTask) class PanopticDeeplabTask(base_task.Task): """A task for Panoptic Deeplab.""" def build_model(self): """Builds panoptic deeplab model.""" 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_panoptic_deeplab( input_specs=input_specs, model_config=self.task_config.model, l2_regularizer=l2_regularizer) return model def initialize(self, model: tf.keras.Model): """Loads pretrained checkpoint.""" if not self.task_config.init_checkpoint: return ckpt_dir_or_file = self.task_config.init_checkpoint if tf.io.gfile.isdir(ckpt_dir_or_file): ckpt_dir_or_file = tf.train.latest_checkpoint(ckpt_dir_or_file) # Restoring checkpoint. if 'all' in self.task_config.init_checkpoint_modules: ckpt = tf.train.Checkpoint(**model.checkpoint_items) status = ckpt.read(ckpt_dir_or_file) status.expect_partial().assert_existing_objects_matched() else: ckpt_items = {} if 'backbone' in self.task_config.init_checkpoint_modules: ckpt_items.update(backbone=model.backbone) if 'decoder' in self.task_config.init_checkpoint_modules: ckpt_items.update(semantic_decoder=model.semantic_decoder) if not self.task_config.model.shared_decoder: ckpt_items.update(instance_decoder=model.instance_decoder) ckpt = tf.train.Checkpoint(**ckpt_items) status = ckpt.read(ckpt_dir_or_file) status.expect_partial().assert_existing_objects_matched() logging.info('Finished loading pretrained checkpoint from %s', ckpt_dir_or_file) def build_inputs(self, params: exp_cfg.DataConfig, input_context: Optional[tf.distribute.InputContext] = None): """Builds panoptic deeplab input.""" decoder_cfg = params.decoder.get() if params.decoder.type == 'simple_decoder': decoder = panoptic_deeplab_input.TfExampleDecoder( regenerate_source_id=decoder_cfg.regenerate_source_id, 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_deeplab_input.Parser( output_size=self.task_config.model.input_size[:2], ignore_label=params.parser.ignore_label, resize_eval_groundtruth=params.parser.resize_eval_groundtruth, groundtruth_padded_size=params.parser.groundtruth_padded_size, aug_scale_min=params.parser.aug_scale_min, aug_scale_max=params.parser.aug_scale_max, aug_rand_hflip=params.parser.aug_rand_hflip, aug_type=params.parser.aug_type, sigma=params.parser.sigma, dtype=params.parser.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, labels: Mapping[str, tf.Tensor], model_outputs: Mapping[str, tf.Tensor], aux_losses: Optional[Any] = None): """Panoptic deeplab losses. Args: labels: labels. model_outputs: Output logits from panoptic deeplab. aux_losses: auxiliarly loss tensors, i.e. `losses` in keras.Model. Returns: The total loss tensor. """ loss_config = self._task_config.losses segmentation_loss_fn = ( panoptic_deeplab_losses.WeightedBootstrappedCrossEntropyLoss( loss_config.label_smoothing, loss_config.class_weights, loss_config.ignore_label, top_k_percent_pixels=loss_config.top_k_percent_pixels)) instance_center_heatmap_loss_fn = panoptic_deeplab_losses.CenterHeatmapLoss( ) instance_center_offset_loss_fn = panoptic_deeplab_losses.CenterOffsetLoss() semantic_weights = tf.cast( labels['semantic_weights'], dtype=model_outputs['instance_centers_heatmap'].dtype) things_mask = tf.cast( tf.squeeze(labels['things_mask'], axis=3), dtype=model_outputs['instance_centers_heatmap'].dtype) valid_mask = tf.cast( tf.squeeze(labels['valid_mask'], axis=3), dtype=model_outputs['instance_centers_heatmap'].dtype) segmentation_loss = segmentation_loss_fn( model_outputs['segmentation_outputs'], labels['category_mask'], sample_weight=semantic_weights) instance_center_heatmap_loss = instance_center_heatmap_loss_fn( model_outputs['instance_centers_heatmap'], labels['instance_centers_heatmap'], sample_weight=valid_mask) instance_center_offset_loss = instance_center_offset_loss_fn( model_outputs['instance_centers_offset'], labels['instance_centers_offset'], sample_weight=things_mask) model_loss = ( loss_config.segmentation_loss_weight * segmentation_loss + loss_config.center_heatmap_loss_weight * instance_center_heatmap_loss + loss_config.center_offset_loss_weight * instance_center_offset_loss) total_loss = model_loss if aux_losses: total_loss += tf.add_n(aux_losses) losses = { 'total_loss': total_loss, 'model_loss': model_loss, 'segmentation_loss': segmentation_loss, 'instance_center_heatmap_loss': instance_center_heatmap_loss, 'instance_center_offset_loss': instance_center_offset_loss } return losses def build_metrics(self, training: bool = True) -> List[ tf.keras.metrics.Metric]: """Build metrics.""" eval_config = self.task_config.evaluation metrics = [] if training: metric_names = [ 'total_loss', 'segmentation_loss', 'instance_center_heatmap_loss', 'instance_center_offset_loss', 'model_loss'] for name in metric_names: metrics.append(tf.keras.metrics.Mean(name, dtype=tf.float32)) if eval_config.report_train_mean_iou: self.train_mean_iou = segmentation_metrics.MeanIoU( name='train_mean_iou', num_classes=self.task_config.model.num_classes, rescale_predictions=False, dtype=tf.float32) else: rescale_predictions = (not self.task_config.validation_data.parser .resize_eval_groundtruth) self.perclass_iou_metric = segmentation_metrics.PerClassIoU( name='per_class_iou', num_classes=self.task_config.model.num_classes, rescale_predictions=rescale_predictions, dtype=tf.float32) if self.task_config.model.generate_panoptic_masks: self.panoptic_quality_metric = ( panoptic_quality_evaluator.PanopticQualityEvaluator( num_categories=self.task_config.model.num_classes, ignored_label=eval_config.ignored_label, max_instances_per_category=eval_config .max_instances_per_category, offset=eval_config.offset, is_thing=eval_config.is_thing, rescale_predictions=eval_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( inputs=images, image_info=labels['image_info'], training=True) outputs = tf.nest.map_structure( lambda x: tf.cast(x, tf.float32), outputs) # Computes per-replica loss. losses = self.build_losses( labels=labels, model_outputs=outputs, aux_losses=model.losses) 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]) if self.task_config.evaluation.report_train_mean_iou: segmentation_labels = { 'masks': labels['category_mask'], 'valid_masks': labels['valid_mask'], 'image_info': labels['image_info'] } self.process_metrics( metrics=[self.train_mean_iou], labels=segmentation_labels, model_outputs=outputs['segmentation_outputs']) logs.update({ self.train_mean_iou.name: self.train_mean_iou.result() }) 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( inputs=images, image_info=labels['image_info'], training=False) logs = {self.loss: 0} segmentation_labels = { 'masks': labels['category_mask'], 'valid_masks': labels['valid_mask'], 'image_info': labels['image_info'] } self.perclass_iou_metric.update_state(segmentation_labels, outputs['segmentation_outputs']) if self.task_config.model.generate_panoptic_masks: pq_metric_labels = { 'category_mask': tf.squeeze(labels['category_mask'], axis=3), 'instance_mask': tf.squeeze(labels['instance_mask'], axis=3), 'image_info': labels['image_info'] } panoptic_outputs = { 'category_mask': outputs['category_mask'], 'instance_mask': outputs['instance_mask'], } logs.update({ self.panoptic_quality_metric.name: (pq_metric_labels, panoptic_outputs)}) return logs def aggregate_logs(self, state=None, step_outputs=None): if state is None: self.perclass_iou_metric.reset_states() state = [self.perclass_iou_metric] if self.task_config.model.generate_panoptic_masks: state += [self.panoptic_quality_metric] if self.task_config.model.generate_panoptic_masks: 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 = {} ious = self.perclass_iou_metric.result() if self.task_config.evaluation.report_per_class_iou: for i, value in enumerate(ious.numpy()): result.update({'segmentation_iou/class_{}'.format(i): value}) # Computes mean IoU result.update({'segmentation_mean_iou': tf.reduce_mean(ious).numpy()}) if self.task_config.model.generate_panoptic_masks: panoptic_quality_results = self.panoptic_quality_metric.result() for k, value in panoptic_quality_results.items(): if k.endswith('per_class'): if self.task_config.evaluation.report_per_class_pq: 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