karolmajek/maxdeeplab
0
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"""This file contains code to build MaX-DeepLab output heads.17 18Reference:19 MaX-DeepLab: "End-to-End Panoptic Segmentation with Mask Transformers",20 CVPR 2021. https://arxiv.org/abs/2012.0075921 Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, Liang-Chieh Chen.22"""23import math24 25import tensorflow as tf26 27from deeplab2 import common28from deeplab2.model.decoder import panoptic_deeplab29from deeplab2.model.layers import convolutions30 31_PIXEL_SPACE_FEATURE_KEY = 'pixel_space_feature'32 33 34def _get_transformer_class_head_num_classes(35 auxiliary_semantic_head_output_channels,36 ignore_label):37 """Computes the num of classes for the transformer class head.38 39 The transformer class head predicts non-void classes (i.e., thing classes and40 stuff classes) and a void (i.e., ∅, no object) class. If the auxiliary41 semantic head output channel includes the void class, e.g., on COCO, we42 directly use the semantic output channel. Otherwise, e.g., on Cityscapes, we43 add 1 (the void class) to the transformer class head.44 45 Args:46 auxiliary_semantic_head_output_channels: An integer, the number of output47 channels of the auxiliary semantic head (it should be the same as the48 num_classes field of the dataset information).49 ignore_label: An integer specifying the ignore label. Default to 255.50 51 Returns:52 num_classes: An integer, the num of classes for the transformer class head.53 """54 if ignore_label >= auxiliary_semantic_head_output_channels:55 return auxiliary_semantic_head_output_channels + 156 else:57 return auxiliary_semantic_head_output_channels58 59 60def add_bias_towards_void(transformer_class_logits, void_prior_prob=0.9):61 """Adds init bias towards the void (no object) class to the class logits.62 63 We initialize the void class with a large probability, similar to Section 3.364 of the Focal Loss paper.65 66 Reference:67 Focal Loss for Dense Object Detection, ICCV 2017.68 https://arxiv.org/abs/1708.0200269 Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, Piotr Dollár.70 71 Args:72 transformer_class_logits: A [batch, num_mask_slots, num_classes] tensor, the73 class logits predicted by the transformer. It concats (num_classes - 1)74 non-void classes, including both thing classes and stuff classes, and the75 void class (the last channel). If the dataset class IDs do not follow this76 order, MaX-DeepLab loss functions will handle the mapping and thus the77 architecture still supports any dataset.78 void_prior_prob: A float, the desired probability (after softmax) of the79 void class at initialization. Defaults to 0.9 as in MaX-DeepLab.80 81 Returns:82 updated_transformer_class_logits: A [batch, num_mask_slots, num_classes]83 84 Raises:85 ValueError: If the rank of transformer_class_logits is not 3.86 """87 class_logits_shape = transformer_class_logits.get_shape().as_list()88 if len(class_logits_shape) != 3:89 raise ValueError('Input transformer_class_logits should have rank 3.')90 91 init_bias = [0.0] * class_logits_shape[-1]92 init_bias[-1] = math.log(93 (class_logits_shape[-1] - 1) * void_prior_prob / (1 - void_prior_prob))94 95 # Broadcasting the 1D init_bias to the 3D transformer_class_logits.96 return transformer_class_logits + tf.constant(init_bias, dtype=tf.float32)97 98 99def batch_norm_on_an_extra_axis(inputs, bn_layer):100 """Applies a batch norm layer on an extra axis.101 102 This batch norm will be used on the pixel space mask logits in MaX-DeepLab to103 avoid careful initialization of previous layers and careful scaling of the104 resulting outputs. In addition, applying batch norm on an extra axis does not105 introduce an extra gamma and beta for each mask slot. Instead, the current106 gamma and beta are shared for all mask slots and do not introduce biases on107 mask slots.108 109 Args:110 inputs: A [batch, height, width, num_mask_slots] tensor.111 bn_layer: A batch norm tf.keras.layers.Layer on the last axis.112 113 Returns:114 outputs: A [batch, height, width, num_mask_slots] tensor.115 """116 expanded_inputs = tf.expand_dims(inputs, axis=-1)117 outputs = bn_layer(expanded_inputs)118 return tf.squeeze(outputs, axis=-1)119 120 121class MaXDeepLab(tf.keras.layers.Layer):122 """A MaX-DeepLab head layer."""123 124 def __init__(self,125 decoder_options,126 max_deeplab_options,127 ignore_label,128 bn_layer=tf.keras.layers.BatchNormalization):129 """Initializes a MaX-DeepLab head.130 131 Args:132 decoder_options: Decoder options as defined in config_pb2.DecoderOptions.133 max_deeplab_options: Model options as defined in134 config_pb2.ModelOptions.MaXDeepLabOptions.135 ignore_label: An integer specifying the ignore label.136 bn_layer: An optional tf.keras.layers.Layer that computes the137 normalization (default: tf.keras.layers.BatchNormalization).138 """139 super(MaXDeepLab, self).__init__(name='MaXDeepLab')140 141 low_level_feature_keys = [142 item.feature_key for item in max_deeplab_options.auxiliary_low_level143 ]144 low_level_channels_project = [145 item.channels_project146 for item in max_deeplab_options.auxiliary_low_level147 ]148 149 self._auxiliary_semantic_decoder = (150 panoptic_deeplab.PanopticDeepLabSingleDecoder(151 high_level_feature_name=decoder_options.feature_key,152 low_level_feature_names=low_level_feature_keys,153 low_level_channels_project=low_level_channels_project,154 aspp_output_channels=decoder_options.aspp_channels,155 decoder_output_channels=decoder_options.decoder_channels,156 atrous_rates=decoder_options.atrous_rates,157 name='auxiliary_semantic_decoder',158 aspp_use_only_1x1_proj_conv=decoder_options159 .aspp_use_only_1x1_proj_conv,160 decoder_conv_type=decoder_options.decoder_conv_type,161 bn_layer=bn_layer))162 self._auxiliary_semantic_head = panoptic_deeplab.PanopticDeepLabSingleHead(163 max_deeplab_options.auxiliary_semantic_head.head_channels,164 max_deeplab_options.auxiliary_semantic_head.output_channels,165 common.PRED_SEMANTIC_LOGITS_KEY,166 name='auxiliary_semantic_head',167 conv_type=max_deeplab_options.auxiliary_semantic_head.head_conv_type,168 bn_layer=bn_layer)169 self._pixel_space_head = panoptic_deeplab.PanopticDeepLabSingleHead(170 max_deeplab_options.pixel_space_head.head_channels,171 max_deeplab_options.pixel_space_head.output_channels,172 _PIXEL_SPACE_FEATURE_KEY,173 name='pixel_space_head',174 conv_type=max_deeplab_options.pixel_space_head.head_conv_type,175 bn_layer=bn_layer)176 177 self._transformer_mask_head = convolutions.Conv1D(178 output_channels=max_deeplab_options.pixel_space_head.output_channels,179 name='transformer_mask_head',180 use_bias=False,181 # Use bn to avoid careful initialization.182 use_bn=True,183 bn_layer=bn_layer,184 bn_gamma_initializer='ones',185 activation=None,186 kernel_initializer='he_normal',187 kernel_size=1,188 padding='valid')189 # The transformer class head predicts non-void classes (i.e., thing classes190 # and stuff classes) and a void (i.e., ∅, no object) class.191 num_classes = _get_transformer_class_head_num_classes(192 max_deeplab_options.auxiliary_semantic_head.output_channels,193 ignore_label=ignore_label)194 self._transformer_class_head = convolutions.Conv1D(195 output_channels=num_classes,196 name='transformer_class_head',197 # Use conv bias rather than bn on this final class logit output.198 use_bias=True,199 use_bn=False,200 activation=None,201 # Follow common ImageNet class initlization with stddev 0.01.202 kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=0.01),203 kernel_size=1,204 padding='valid')205 206 self._pixel_space_feature_batch_norm = bn_layer(207 axis=-1, name='pixel_space_feature_batch_norm',208 gamma_initializer=tf.keras.initializers.Constant(1.0))209 # Use a batch norm to avoid care initialization of the mask outputs.210 self._pixel_space_mask_batch_norm = bn_layer(211 axis=-1, name='pixel_space_mask_batch_norm',212 # Initialize the pixel space mask with a low temperature.213 gamma_initializer=tf.keras.initializers.Constant(0.1))214 215 def reset_pooling_layer(self):216 """Resets the ASPP pooling layers to global average pooling."""217 self._auxiliary_semantic_decoder.reset_pooling_layer()218 219 def set_pool_size(self, pool_size):220 """Sets the pooling size of the ASPP pooling layers.221 222 Args:223 pool_size: A tuple specifying the pooling size of the ASPP pooling layers.224 """225 self._auxiliary_semantic_decoder.set_pool_size(pool_size)226 227 def get_pool_size(self):228 return self._auxiliary_semantic_decoder.get_pool_size()229 230 @property231 def checkpoint_items(self):232 items = {233 common.CKPT_SEMANTIC_DECODER:234 self._auxiliary_semantic_decoder,235 common.CKPT_SEMANTIC_HEAD_WITHOUT_LAST_LAYER:236 self._auxiliary_semantic_head.conv_block,237 common.CKPT_SEMANTIC_LAST_LAYER:238 self._auxiliary_semantic_head.final_conv,239 common.CKPT_PIXEL_SPACE_HEAD:240 self._pixel_space_head,241 common.CKPT_TRANSFORMER_MASK_HEAD:242 self._transformer_mask_head,243 common.CKPT_TRANSFORMER_CLASS_HEAD:244 self._transformer_class_head,245 common.CKPT_PIXEL_SPACE_FEATURE_BATCH_NORM:246 self._pixel_space_feature_batch_norm,247 common.CKPT_PIXEL_SPACE_MASK_BATCH_NORM:248 self._pixel_space_mask_batch_norm,249 }250 return items251 252 def call(self, features, training=False):253 """Performs a forward pass.254 255 Args:256 features: An input dict of tf.Tensor with shape [batch, height, width,257 channels] or [batch, length, channels]. Different keys should point to258 different features extracted by the encoder, e.g., low-level or259 high-level features.260 training: A boolean flag indicating whether training behavior should be261 used (default: False).262 263 Returns:264 A dictionary containing the auxiliary semantic segmentation logits, the265 pixel space normalized feature, the pixel space mask logits, and the266 mask transformer class logits.267 """268 results = {}269 semantic_features = features['feature_semantic']270 panoptic_features = features['feature_panoptic']271 transformer_class_feature = features['transformer_class_feature']272 transformer_mask_feature = features['transformer_mask_feature']273 274 # Auxiliary semantic head.275 semantic_shape = semantic_features.get_shape().as_list()276 panoptic_shape = panoptic_features.get_shape().as_list()277 # MaX-DeepLab always predicts panoptic feature at high resolution (e.g.,278 # stride 4 or stride 2), but the auxiliary semantic feature could be at low279 # resolution (e.g., stride 16 or stride 32), in the absence of the stacked280 # decoder (L == 0). In this case, we use an auxiliary semantic decoder on281 # top of the semantic feature, in order to add the auxiliary semantic loss.282 if semantic_shape[1:3] != panoptic_shape[1:3]:283 semantic_features = self._auxiliary_semantic_decoder(284 features, training=training)285 auxiliary_semantic_results = self._auxiliary_semantic_head(286 semantic_features, training=training)287 results.update(auxiliary_semantic_results)288 289 # Pixel space head.290 pixel_space_feature = self._pixel_space_head(291 panoptic_features, training=training)[_PIXEL_SPACE_FEATURE_KEY]292 pixel_space_feature = self._pixel_space_feature_batch_norm(293 pixel_space_feature)294 pixel_space_normalized_feature = tf.math.l2_normalize(295 pixel_space_feature, axis=-1)296 results[common.PRED_PIXEL_SPACE_NORMALIZED_FEATURE_KEY] = (297 pixel_space_normalized_feature)298 299 # Transformer class head.300 transformer_class_logits = self._transformer_class_head(301 transformer_class_feature)302 # Bias towards the void class at initialization.303 transformer_class_logits = add_bias_towards_void(304 transformer_class_logits)305 results[common.PRED_TRANSFORMER_CLASS_LOGITS_KEY] = transformer_class_logits306 307 # Transformer mask kernel.308 transformer_mask_kernel = self._transformer_mask_head(309 transformer_mask_feature)310 311 # Convolutional mask head. The pixel space mask logits are the matrix312 # multiplication (or convolution) of the pixel space normalized feature and313 # the transformer mask kernel.314 pixel_space_mask_logits = tf.einsum(315 'bhwd,bid->bhwi',316 pixel_space_normalized_feature,317 transformer_mask_kernel)318 # The above multiplication constructs a second-order operation which is319 # sensitive to the feature scales and initializations. In order to avoid320 # careful initialization or scaling of the layers, we apply batch norms on321 # top of pixel_space_feature, transformer_mask_kernel, and the resulting322 # pixel_space_mask_logits.323 pixel_space_mask_logits = batch_norm_on_an_extra_axis(324 pixel_space_mask_logits, self._pixel_space_mask_batch_norm)325 results[common.PRED_PIXEL_SPACE_MASK_LOGITS_KEY] = (326 pixel_space_mask_logits)327 328 return results329