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 an ASPP layer.17 18Reference:19 - [Rethinking Atrous Convolution for Semantic Image Segmentation](20 https://arxiv.org/pdf/1706.05587.pdf)21 - [ParseNet: Looking Wider to See Better](22 https://arxiv.org/pdf/1506.04579.pdf).23"""24from absl import logging25import tensorflow as tf26 27from deeplab2.model import utils28from deeplab2.model.layers import convolutions29 30 31layers = tf.keras.layers32backend = tf.keras.backend33 34 35class ASPPConv(tf.keras.layers.Layer):36 """An atrous convolution for ASPP."""37 38 def __init__(self,39 output_channels,40 atrous_rate,41 name,42 bn_layer=tf.keras.layers.BatchNormalization):43 """Creates a atrous convolution layer for the ASPP.44 45 This layer consists of an atrous convolution followed by a BatchNorm layer46 and a ReLU activation.47 48 Args:49 output_channels: An integer specifying the number of output channels of50 the convolution.51 atrous_rate: An integer specifying the atrous/dilation rate of the52 convolution.53 name: A string specifying the name of this layer.54 bn_layer: An optional tf.keras.layers.Layer that computes the55 normalization (default: tf.keras.layers.BatchNormalization).56 """57 super(ASPPConv, self).__init__(name=name)58 59 self._conv_bn_act = convolutions.Conv2DSame(60 output_channels,61 kernel_size=3,62 name='conv_bn_act',63 atrous_rate=atrous_rate,64 use_bias=False,65 use_bn=True,66 bn_layer=bn_layer,67 activation='relu')68 69 def call(self, input_tensor, training=False):70 """Performs a forward pass.71 72 Args:73 input_tensor: An input tensor of type tf.Tensor with shape [batch, height,74 width, channels].75 training: A boolean flag indicating whether training behavior should be76 used (default: False).77 78 Returns:79 The output tensor.80 """81 return self._conv_bn_act(input_tensor, training=training)82 83 84class ASPPPool(tf.keras.layers.Layer):85 """A pooling layer for ASPP."""86 87 def __init__(self,88 output_channels,89 name,90 bn_layer=tf.keras.layers.BatchNormalization):91 """Creates a pooling layer for the ASPP.92 93 This layer consists of a global average pooling, followed by a convolution,94 and by a BatchNorm layer and a ReLU activation.95 96 Args:97 output_channels: An integer specifying the number of output channels of98 the convolution.99 name: A string specifying the name of this layer.100 bn_layer: An optional tf.keras.layers.Layer that computes the101 normalization (default: tf.keras.layers.BatchNormalization).102 """103 super(ASPPPool, self).__init__(name=name)104 105 self._pool_size = (None, None)106 self._conv_bn_act = convolutions.Conv2DSame(107 output_channels,108 kernel_size=1,109 name='conv_bn_act',110 use_bias=False,111 use_bn=True,112 bn_layer=bn_layer,113 activation='relu')114 115 def set_pool_size(self, pool_size):116 """Sets the pooling size of the pooling layer.117 118 The default behavior of the pooling layer is global average pooling. A119 custom pooling size can be set here.120 121 Args:122 pool_size: A tuple specifying the pooling size of the pooling layer.123 124 Raises:125 An error occurs if exactly one pooling dimension is set to 'None'.126 """127 # If exactly one pooling dimension is 'None' raise an error.128 if None in pool_size and pool_size != (None, None):129 raise ValueError('The ASPP pooling layer requires that the pooling size '130 'is set explicitly for both dimensions. In case, global '131 'average pooling should be used, call '132 'reset_pooling_layer() or set both to None.')133 134 self._pool_size = pool_size135 logging.info('Global average pooling in the ASPP pooling layer was replaced'136 ' with tiled average pooling using the provided pool_size. '137 'Please make sure this behavior is intended.')138 139 def get_pool_size(self):140 return self._pool_size141 142 def reset_pooling_layer(self):143 """Resets the pooling layer to global average pooling."""144 self._pool_size = (None, None)145 146 def call(self, input_tensor, training=False):147 """Performs a forward pass.148 149 Args:150 input_tensor: An input tensor of type tf.Tensor with shape [batch, height,151 width, channels].152 training: A boolean flag indicating whether training behavior should be153 used (default: False).154 155 Returns:156 The output tensor.157 """158 if tuple(self._pool_size) == (None, None):159 # Global image pooling160 pool_size = input_tensor.shape[1:3]161 else:162 # Tiled image pooling163 pool_size = self._pool_size164 165 x = backend.pool2d(input_tensor, pool_size, padding='valid',166 pool_mode='avg')167 x = self._conv_bn_act(x, training=training)168 169 target_h = tf.shape(input_tensor)[1]170 target_w = tf.shape(input_tensor)[2]171 172 x = utils.resize_align_corners(x, [target_h, target_w])173 return x174 175 176class ASPP(tf.keras.layers.Layer):177 """An atrous spatial pyramid pooling layer."""178 179 def __init__(self,180 output_channels,181 atrous_rates,182 aspp_use_only_1x1_proj_conv=False,183 name='ASPP',184 bn_layer=tf.keras.layers.BatchNormalization):185 """Creates an ASPP layer.186 187 Args:188 output_channels: An integer specifying the number of output channels of189 each ASPP convolution layer.190 atrous_rates: A list of three integers specifying the atrous/dilation rate191 of each ASPP convolution layer.192 aspp_use_only_1x1_proj_conv: Boolean, specifying if the ASPP five branches193 are turned off or not. If True, the ASPP module is degenerated to one194 1x1 convolution, projecting the input channels to `output_channels`.195 name: A string specifying the name of this layer (default: 'ASPP').196 bn_layer: An optional tf.keras.layers.Layer that computes the197 normalization (default: tf.keras.layers.BatchNormalization).198 199 Raises:200 ValueError: An error occurs when both atrous_rates does not contain 3201 elements and `aspp_use_only_1x1_proj_conv` is False.202 """203 super(ASPP, self).__init__(name=name)204 205 if not aspp_use_only_1x1_proj_conv and len(atrous_rates) != 3:206 raise ValueError(207 'The ASPP layers need exactly 3 atrous rates, but %d were given' %208 len(atrous_rates))209 self._aspp_use_only_1x1_proj_conv = aspp_use_only_1x1_proj_conv210 211 # Projection convolution is always used.212 self._proj_conv_bn_act = convolutions.Conv2DSame(213 output_channels,214 kernel_size=1,215 name='proj_conv_bn_act',216 use_bias=False,217 use_bn=True,218 bn_layer=bn_layer,219 activation='relu')220 221 if not aspp_use_only_1x1_proj_conv:222 self._conv_bn_act = convolutions.Conv2DSame(223 output_channels,224 kernel_size=1,225 name='conv_bn_act',226 use_bias=False,227 use_bn=True,228 bn_layer=bn_layer,229 activation='relu')230 rate1, rate2, rate3 = atrous_rates231 self._aspp_conv1 = ASPPConv(output_channels, rate1, name='aspp_conv1',232 bn_layer=bn_layer)233 self._aspp_conv2 = ASPPConv(output_channels, rate2, name='aspp_conv2',234 bn_layer=bn_layer)235 self._aspp_conv3 = ASPPConv(output_channels, rate3, name='aspp_conv3',236 bn_layer=bn_layer)237 self._aspp_pool = ASPPPool(output_channels, name='aspp_pool',238 bn_layer=bn_layer)239 # Dropout is needed only when ASPP five branches are used.240 self._proj_drop = layers.Dropout(rate=0.1)241 242 def set_pool_size(self, pool_size):243 """Sets the pooling size of the ASPP pooling layer.244 245 The default behavior of the pooling layer is global average pooling. A246 custom pooling size can be set here.247 248 Args:249 pool_size: A tuple specifying the pooling size of the ASPP pooling layer.250 """251 if not self._aspp_use_only_1x1_proj_conv:252 self._aspp_pool.set_pool_size(pool_size)253 254 def get_pool_size(self):255 if not self._aspp_use_only_1x1_proj_conv:256 return self._aspp_pool.get_pool_size()257 else:258 return (None, None)259 260 def reset_pooling_layer(self):261 """Resets the pooling layer to global average pooling."""262 self._aspp_pool.reset_pooling_layer()263 264 def call(self, input_tensor, training=False):265 """Performs a forward pass.266 267 Args:268 input_tensor: An input tensor of type tf.Tensor with shape [batch, height,269 width, channels].270 training: A boolean flag indicating whether training behavior should be271 used (default: False).272 273 Returns:274 The output tensor.275 """276 if self._aspp_use_only_1x1_proj_conv:277 x = self._proj_conv_bn_act(input_tensor, training=training)278 else:279 # Apply the ASPP module.280 results = []281 results.append(self._conv_bn_act(input_tensor, training=training))282 results.append(self._aspp_conv1(input_tensor, training=training))283 results.append(self._aspp_conv2(input_tensor, training=training))284 results.append(self._aspp_conv3(input_tensor, training=training))285 results.append(self._aspp_pool(input_tensor, training=training))286 x = tf.concat(results, 3)287 x = self._proj_conv_bn_act(x, training=training)288 x = self._proj_drop(x, training=training)289 return x290