def build(self, input_shape): height = input_shape[1] width = input_shape[2] channels = input_shape[3] norm_layer = (tf.keras.layers.experimental.SyncBatchNormalization if self._use_sync_bn else tf.keras.layers.BatchNormalization) norm_with_quantize = helper.BatchNormalizationQuantized(norm_layer) norm_no_quantize = helper.BatchNormalizationNoQuantized(norm_layer) norm = helper.norm_by_activation(self._activation, norm_with_quantize, norm_no_quantize) self.aspp_layers = [] conv1 = helper.Conv2DQuantized( filters=self._output_channels, kernel_size=(1, 1), kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_bias=False, activation=helper.NoOpActivation()) norm1 = norm(axis=self._bn_axis, momentum=self._batchnorm_momentum, epsilon=self._batchnorm_epsilon) self.aspp_layers.append([conv1, norm1]) for dilation_rate in self._dilation_rates: leading_layers = [] kernel_size = (3, 3) if self._use_depthwise_convolution: leading_layers += [ helper.DepthwiseConv2DOutputQuantized( depth_multiplier=1, kernel_size=kernel_size, padding='same', depthwise_regularizer=self._kernel_regularizer, depthwise_initializer=self._kernel_initializer, dilation_rate=dilation_rate, use_bias=False, activation=helper.NoOpActivation()) ] kernel_size = (1, 1) conv_dilation = leading_layers + [ helper.Conv2DQuantized( filters=self._output_channels, kernel_size=kernel_size, padding='same', kernel_regularizer=self._kernel_regularizer, kernel_initializer=self._kernel_initializer, dilation_rate=dilation_rate, use_bias=False, activation=helper.NoOpActivation()) ] norm_dilation = norm(axis=self._bn_axis, momentum=self._batchnorm_momentum, epsilon=self._batchnorm_epsilon) self.aspp_layers.append(conv_dilation + [norm_dilation]) if self._pool_kernel_size is None: pooling = [ helper.GlobalAveragePooling2DQuantized(), helper.ReshapeQuantized((1, 1, channels)) ] else: pooling = [ helper.AveragePooling2DQuantized(self._pool_kernel_size) ] conv2 = helper.Conv2DQuantized( filters=self._output_channels, kernel_size=(1, 1), kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_bias=False, activation=helper.NoOpActivation()) norm2 = norm(axis=self._bn_axis, momentum=self._batchnorm_momentum, epsilon=self._batchnorm_epsilon) self.aspp_layers.append(pooling + [conv2, norm2]) self._resizing_layer = helper.ResizingQuantized( height, width, interpolation=self._interpolation) self._projection = [ helper.Conv2DQuantized(filters=self._output_channels, kernel_size=(1, 1), kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_bias=False, activation=helper.NoOpActivation()), norm_with_quantize(axis=self._bn_axis, momentum=self._batchnorm_momentum, epsilon=self._batchnorm_epsilon) ] self._dropout_layer = tf.keras.layers.Dropout(rate=self._dropout) self._concat_layer = helper.ConcatenateQuantized(axis=-1)
def build(self, input_shape: Sequence[tf.TensorShape]): """Creates the variables of the segmentation head.""" # When input_shape is a list/tuple, the first corresponds to backbone # features used for resizing the decoder features (the second) if feature # fusion type is `deeplabv3plus`. backbone_shape = input_shape[0] use_depthwise_convolution = self._config_dict[ 'use_depthwise_convolution'] random_initializer = tf.keras.initializers.RandomNormal(stddev=0.01) conv_kwargs = { 'kernel_size': 3 if not use_depthwise_convolution else 1, 'padding': 'same', 'use_bias': False, 'kernel_initializer': random_initializer, 'kernel_regularizer': self._config_dict['kernel_regularizer'], } norm_layer = (tf.keras.layers.experimental.SyncBatchNormalization if self._config_dict['use_sync_bn'] else tf.keras.layers.BatchNormalization) norm_with_quantize = helper.BatchNormalizationQuantized(norm_layer) norm_no_quantize = helper.BatchNormalizationNoQuantized(norm_layer) norm = helper.norm_by_activation(self._config_dict['activation'], norm_with_quantize, norm_no_quantize) bn_kwargs = { 'axis': self._bn_axis, 'momentum': self._config_dict['norm_momentum'], 'epsilon': self._config_dict['norm_epsilon'], } if self._config_dict['feature_fusion'] == 'deeplabv3plus': # Deeplabv3+ feature fusion layers. self._dlv3p_conv = helper.Conv2DQuantized( kernel_size=1, padding='same', use_bias=False, kernel_initializer=tf.keras.initializers.RandomNormal( stddev=0.01), kernel_regularizer=self._config_dict['kernel_regularizer'], name='segmentation_head_deeplabv3p_fusion_conv', filters=self._config_dict['low_level_num_filters'], activation=helper.NoOpActivation()) self._dlv3p_norm = norm( name='segmentation_head_deeplabv3p_fusion_norm', **bn_kwargs) # Segmentation head layers. self._convs = [] self._norms = [] for i in range(self._config_dict['num_convs']): if use_depthwise_convolution: self._convs.append( helper.DepthwiseConv2DQuantized( name='segmentation_head_depthwise_conv_{}'.format(i), kernel_size=3, padding='same', use_bias=False, depthwise_initializer=random_initializer, depthwise_regularizer=self. _config_dict['kernel_regularizer'], depth_multiplier=1, activation=helper.NoOpActivation())) norm_name = 'segmentation_head_depthwise_norm_{}'.format(i) self._norms.append(norm(name=norm_name, **bn_kwargs)) conv_name = 'segmentation_head_conv_{}'.format(i) self._convs.append( helper.Conv2DQuantized( name=conv_name, filters=self._config_dict['num_filters'], activation=helper.NoOpActivation(), **conv_kwargs)) norm_name = 'segmentation_head_norm_{}'.format(i) self._norms.append(norm(name=norm_name, **bn_kwargs)) self._classifier = helper.Conv2DOutputQuantized( name='segmentation_output', filters=self._config_dict['num_classes'], kernel_size=self._config_dict['prediction_kernel_size'], padding='same', bias_initializer=tf.zeros_initializer(), kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.01), kernel_regularizer=self._config_dict['kernel_regularizer'], bias_regularizer=self._config_dict['bias_regularizer'], activation=helper.NoOpActivation()) self._upsampling_layer = helper.UpSampling2DQuantized( size=(self._config_dict['upsample_factor'], self._config_dict['upsample_factor']), interpolation='nearest') self._resizing_layer = tf.keras.layers.Resizing( backbone_shape[1], backbone_shape[2], interpolation='bilinear') self._concat_layer = helper.ConcatenateQuantized(axis=self._bn_axis) super().build(input_shape)