karolmajek/Axial-DeepLab-SWideRNet

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""Utility script to perform net surgery on a model.17 18This script will perform net surgery on DeepLab models trained on a source19dataset and create a new checkpoint for the target dataset.20"""21 22from typing import Any, Dict, Text, Tuple23 24from absl import app25from absl import flags26from absl import logging27 28import numpy as np29import tensorflow as tf30 31from google.protobuf import text_format32from deeplab2 import common33from deeplab2 import config_pb234from deeplab2.data import dataset35from deeplab2.model import deeplab36 37FLAGS = flags.FLAGS38 39flags.DEFINE_string('source_dataset', 'cityscapes',40                    'Dataset name on which the model has been pretrained. '41                    'Supported datasets: `cityscapes`.')42 43flags.DEFINE_string('target_dataset', 'motchallenge_step',44                    'Dataset name for conversion. Supported datasets: '45                    '`motchallenge_step`.')46 47flags.DEFINE_string('input_config_path', None,48                    'Path to a config file that defines the DeepLab model and '49                    'the checkpoint path.')50 51flags.DEFINE_string('output_checkpoint_path', None,52                    'Output filename for the generated checkpoint file.')53 54 55_SUPPORTED_SOURCE_DATASETS = {'cityscapes'}56_SUPPORTED_TARGET_DATASETS = {'motchallenge_step'}57 58_CITYSCAPES_TO_MOTCHALLENGE_STEP = (59    1,  # sidewalk60    2,  # building61    8,  # vegetation62    10,  # sky63    11,  # pedestrian64    12,  # rider65    18,  # bicycle66)67 68_DATASET_TO_INFO = {69    'cityscapes': dataset.CITYSCAPES_PANOPTIC_INFORMATION,70    'motchallenge_step': dataset.MOTCHALLENGE_STEP_INFORMATION,71}72_INPUT_SIZE = (1025, 2049, 3)73 74 75def _load_model(76    config_path: Text,77    source_dataset: Text) -> Tuple[deeplab.DeepLab,78                                   config_pb2.ExperimentOptions]:79  """Load DeepLab model based on config and dataset."""80  options = config_pb2.ExperimentOptions()81  with tf.io.gfile.GFile(config_path) as f:82    text_format.Parse(f.read(), options)83  options.model_options.panoptic_deeplab.semantic_head.output_channels = (84      _DATASET_TO_INFO[source_dataset].num_classes)85  model = deeplab.DeepLab(options,86                          _DATASET_TO_INFO[source_dataset])87  return model, options88 89 90def _convert_bias(input_tensor: np.ndarray,91                  label_list: Tuple[int, ...]) -> np.ndarray:92  """Converts 1D tensor bias w.r.t. label list.93 94  We select the subsets from the input_tensor based on the label_list.95 96  We assume input_tensor has shape = [num_classes], where97  input_tensor is the bias weights trained on source dataset, and num_classes98  is the number of classes in source dataset.99 100  Args:101    input_tensor: A numpy array with ndim == 1.102    label_list: A tuple of labels used for net surgery.103 104  Returns:105    A numpy array with values modified.106 107  Raises:108    ValueError: input_tensor's ndim != 1.109  """110  if input_tensor.ndim != 1:111    raise ValueError('The bias tensor should have ndim == 1.')112 113  num_elements = len(label_list)114  output_tensor = np.zeros(num_elements, dtype=np.float32)115  for i, label in enumerate(label_list):116    output_tensor[i] = input_tensor[label]117  return output_tensor118 119 120def _convert_kernels(input_tensor: np.ndarray,121                     label_list: Tuple[int, ...]) -> np.ndarray:122  """Converts 4D tensor kernels w.r.t. label list.123 124  We select the subsets from the input_tensor based on the label_list.125 126  We assume input_tensor has shape = [h, w, input_dim, num_classes], where127  input_tensor is the kernel weights trained on source dataset, and num_classes128  is the number of classes in source dataset.129 130  Args:131    input_tensor: A numpy array with ndim == 4.132    label_list: A tuple of labels used for net surgery.133 134  Returns:135    A numpy array with values modified.136 137  Raises:138    ValueError: input_tensor's ndim != 4.139  """140  if input_tensor.ndim != 4:141    raise ValueError('The kernels tensor should have ndim == 4.')142 143  num_elements = len(label_list)144  kernel_height, kernel_width, input_dim, _ = input_tensor.shape145  output_tensor = np.zeros(146      (kernel_height, kernel_width, input_dim, num_elements), dtype=np.float32)147  for i, label in enumerate(label_list):148    output_tensor[:, :, :, i] = input_tensor[:, :, :, label]149  return output_tensor150 151 152def _restore_checkpoint(restore_dict: Dict[Any, Any],153                        options: config_pb2.ExperimentOptions154                        ) -> tf.train.Checkpoint:155  """Reads the provided dict items from the checkpoint specified in options.156 157  Args:158    restore_dict: A mapping of checkpoint item to location.159    options: A experiment configuration containing the checkpoint location.160 161  Returns:162    The loaded checkpoint.163  """164  ckpt = tf.train.Checkpoint(**restore_dict)165  if tf.io.gfile.isdir(options.model_options.initial_checkpoint):166    path = tf.train.latest_checkpoint(167        options.model_options.initial_checkpoint)168    status = ckpt.restore(path)169  else:170    status = ckpt.restore(options.model_options.initial_checkpoint)171  status.expect_partial().assert_existing_objects_matched()172  return ckpt173 174 175def main(_) -> None:176  if FLAGS.source_dataset not in _SUPPORTED_SOURCE_DATASETS:177    raise ValueError('Source dataset is not supported. Use --help to get list '178                     'of supported datasets.')179  if FLAGS.target_dataset not in _SUPPORTED_TARGET_DATASETS:180    raise ValueError('Target dataset is not supported. Use --help to get list '181                     'of supported datasets.')182 183  logging.info('Loading DeepLab model from config %s', FLAGS.input_config_path)184  source_model, options = _load_model(FLAGS.input_config_path,185                                      FLAGS.source_dataset)186  logging.info('Load pretrained checkpoint.')187  _restore_checkpoint(source_model.checkpoint_items, options)188  source_model(tf.keras.Input(_INPUT_SIZE), training=False)189 190  logging.info('Perform net surgery.')191  semantic_weights = (192      source_model._decoder._semantic_head.final_conv.get_weights())  # pylint: disable=protected-access193 194  if (FLAGS.source_dataset == 'cityscapes' and195      FLAGS.target_dataset == 'motchallenge_step'):196    # Kernels.197    semantic_weights[0] = _convert_kernels(semantic_weights[0],198                                           _CITYSCAPES_TO_MOTCHALLENGE_STEP)199    # Bias.200    semantic_weights[1] = _convert_bias(semantic_weights[1],201                                        _CITYSCAPES_TO_MOTCHALLENGE_STEP)202 203  logging.info('Load target model without last semantic layer.')204  target_model, _ = _load_model(FLAGS.input_config_path, FLAGS.target_dataset)205  restore_dict = target_model.checkpoint_items206  del restore_dict[common.CKPT_SEMANTIC_LAST_LAYER]207 208  ckpt = _restore_checkpoint(restore_dict, options)209  target_model(tf.keras.Input(_INPUT_SIZE), training=False)210  target_model._decoder._semantic_head.final_conv.set_weights(semantic_weights)  # pylint: disable=protected-access211 212  logging.info('Save checkpoint to output path: %s',213               FLAGS.output_checkpoint_path)214  ckpt = tf.train.Checkpoint(**target_model.checkpoint_items)215  ckpt.save(FLAGS.output_checkpoint_path)216 217 218if __name__ == '__main__':219  flags.mark_flags_as_required(220      ['input_config_path', 'output_checkpoint_path'])221  app.run(main)222