mosaic_config.py

# 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.

"""Configuration definition for Semantic Segmentation with MOSAIC."""
import dataclasses
import os
from typing import List, Optional, Union

import numpy as np
from official.core import config_definitions as cfg
from official.core import exp_factory
from official.modeling import hyperparams
from official.modeling import optimization
from official.vision.configs import common
from official.vision.configs import semantic_segmentation as seg_cfg
from official.vision.configs.google import backbones


@dataclasses.dataclass
class MosaicDecoderHead(hyperparams.Config):
  """MOSAIC decoder head config for Segmentation."""
  num_classes: int = 19
  decoder_input_levels: List[str] = dataclasses.field(default_factory=list)
  decoder_stage_merge_styles: List[str] = dataclasses.field(
      default_factory=list)
  decoder_filters: List[int] = dataclasses.field(default_factory=list)
  decoder_projected_filters: List[int] = dataclasses.field(default_factory=list)
  encoder_end_level: int = 4
  use_additional_classifier_layer: bool = False
  classifier_kernel_size: int = 1
  activation: str = 'relu'
  kernel_initializer: str = 'glorot_uniform'
  interpolation: str = 'bilinear'


@dataclasses.dataclass
class MosaicEncoderNeck(hyperparams.Config):
  """MOSAIC encoder neck config for segmentation."""
  encoder_input_level: Union[str, int] = '4'
  branch_filter_depths: List[int] = dataclasses.field(default_factory=list)
  conv_kernel_sizes: List[int] = dataclasses.field(default_factory=list)
  pyramid_pool_bin_nums: List[int] = dataclasses.field(default_factory=list)
  activation: str = 'relu'
  dropout_rate: float = 0.1
  kernel_initializer: str = 'glorot_uniform'
  interpolation: str = 'bilinear'
  use_depthwise_convolution: bool = True


@dataclasses.dataclass
class MosaicSemanticSegmentationModel(hyperparams.Config):
  """MOSAIC semantic segmentation model config."""
  num_classes: int = 19
  input_size: List[int] = dataclasses.field(default_factory=list)
  head: MosaicDecoderHead = MosaicDecoderHead()
  backbone: backbones.Backbone = backbones.Backbone(
      type='mobilenet', mobilenet=backbones.MobileNet())
  neck: MosaicEncoderNeck = MosaicEncoderNeck()
  norm_activation: common.NormActivation = common.NormActivation(
      use_sync_bn=True, norm_momentum=0.99, norm_epsilon=0.001)


@dataclasses.dataclass
class MosaicSemanticSegmentationTask(seg_cfg.SemanticSegmentationTask):
  """The config for MOSAIC segmentation task."""
  model: MosaicSemanticSegmentationModel = MosaicSemanticSegmentationModel()
  train_data: seg_cfg.DataConfig = seg_cfg.DataConfig(is_training=True)
  validation_data: seg_cfg.DataConfig = seg_cfg.DataConfig(is_training=False)
  losses: seg_cfg.Losses = seg_cfg.Losses()
  evaluation: seg_cfg.Evaluation = seg_cfg.Evaluation()
  train_input_partition_dims: List[int] = dataclasses.field(
      default_factory=list)
  eval_input_partition_dims: List[int] = dataclasses.field(
      default_factory=list)
  init_checkpoint: Optional[str] = None
  init_checkpoint_modules: Union[
      str, List[str]] = 'all'  # all, backbone, and/or neck.
  export_config: seg_cfg.ExportConfig = seg_cfg.ExportConfig()


# Cityscapes Dataset (Download and process the dataset yourself)
CITYSCAPES_TRAIN_EXAMPLES = 2975
CITYSCAPES_VAL_EXAMPLES = 500
CITYSCAPES_INPUT_PATH_BASE = 'cityscapes/tfrecord'


@exp_factory.register_config_factory('mosaic_mnv35_cityscapes')
def mosaic_mnv35_cityscapes() -> cfg.ExperimentConfig:
  """Instantiates an experiment configuration of image segmentation task.

  This image segmentation experiment is conducted on Cityscapes dataset. The
  model architecture is a MOSAIC encoder-decoer. The default backbone network is
  a mobilenet variant called Mobilenet_v3.5-MultiAvg on top of which the MOSAIC
  encoder-decoder can be deployed. All detailed configurations can be overridden
  by a .yaml file provided by the user to launch the experiments. Please refer
  to .yaml examples in the path of ../configs/experiments/.

  Returns:
    A particular instance of cfg.ExperimentConfig for MOSAIC model based
    image semantic segmentation task.
  """
  train_batch_size = 16
  eval_batch_size = 16
  steps_per_epoch = CITYSCAPES_TRAIN_EXAMPLES // train_batch_size
  output_stride = 16

  backbone_output_level = int(np.math.log2(output_stride))
  config = cfg.ExperimentConfig(
      task=MosaicSemanticSegmentationTask(
          model=MosaicSemanticSegmentationModel(
              # Cityscapes uses only 19 semantic classes for train/evaluation.
              # The void (background) class is ignored in train and evaluation.
              num_classes=19,
              input_size=[None, None, 3],
              backbone=backbones.Backbone(
                  type='mobilenet',
                  mobilenet=backbones.MobileNet(
                      model_id='MobileNetMultiAVGSeg',
                      output_intermediate_endpoints=True,
                      output_stride=output_stride)),
              neck=MosaicEncoderNeck(
                  encoder_input_level=backbone_output_level,
                  branch_filter_depths=[64, 64],
                  conv_kernel_sizes=[3, 5],
                  pyramid_pool_bin_nums=[1, 4, 8, 16],  # paper default
                  activation='relu',
                  dropout_rate=0.1,
                  kernel_initializer='glorot_uniform',
                  interpolation='bilinear',
                  use_depthwise_convolution=True),
              head=MosaicDecoderHead(
                  num_classes=19,
                  decoder_input_levels=['3/depthwise', '2/depthwise'],
                  decoder_stage_merge_styles=['concat_merge', 'sum_merge'],
                  decoder_filters=[64, 64],
                  decoder_projected_filters=[19, 19],
                  encoder_end_level=backbone_output_level,
                  use_additional_classifier_layer=False,
                  classifier_kernel_size=1,
                  activation='relu',
                  kernel_initializer='glorot_uniform',
                  interpolation='bilinear'),
              norm_activation=common.NormActivation(
                  activation='relu',
                  norm_momentum=0.99,
                  norm_epsilon=1e-3,
                  use_sync_bn=True)),
          losses=seg_cfg.Losses(l2_weight_decay=4e-5),
          train_data=seg_cfg.DataConfig(
              input_path=os.path.join(CITYSCAPES_INPUT_PATH_BASE,
                                      'train_fine**'),
              crop_size=[1024, 2048],
              output_size=[1024, 2048],
              is_training=True,
              global_batch_size=train_batch_size,
              aug_scale_min=0.5,
              aug_scale_max=2.0),
          validation_data=seg_cfg.DataConfig(
              input_path=os.path.join(CITYSCAPES_INPUT_PATH_BASE, 'val_fine*'),
              output_size=[1024, 2048],
              is_training=False,
              global_batch_size=eval_batch_size,
              resize_eval_groundtruth=True,
              drop_remainder=False),
          # Imagenet pre-trained Mobilenet_v3.5-MultiAvg checkpoint.
          init_checkpoint='gs://tf_model_garden/vision/mobilenet/v3.5multiavg_seg_float/',
          init_checkpoint_modules='backbone'),
      trainer=cfg.TrainerConfig(
          steps_per_loop=steps_per_epoch,
          summary_interval=steps_per_epoch,
          checkpoint_interval=steps_per_epoch,
          train_steps=100000,
          validation_steps=CITYSCAPES_VAL_EXAMPLES // eval_batch_size,
          validation_interval=steps_per_epoch,
          best_checkpoint_eval_metric='mean_iou',
          best_checkpoint_export_subdir='best_ckpt',
          best_checkpoint_metric_comp='higher',
          optimizer_config=optimization.OptimizationConfig({
              'optimizer': {
                  'type': 'sgd',
                  'sgd': {
                      'momentum': 0.9
                  }
              },
              'learning_rate': {
                  'type': 'polynomial',
                  'polynomial': {
                      'initial_learning_rate': 0.1,
                      'decay_steps': 100000,
                      'end_learning_rate': 0.0,
                      'power': 0.9
                  }
              },
              'warmup': {
                  'type': 'linear',
                  'linear': {
                      'warmup_steps': 5 * steps_per_epoch,
                      'warmup_learning_rate': 0
                  }
              }
          })),
      restrictions=[
          'task.train_data.is_training != None',
          'task.validation_data.is_training != None'
      ])

  return config