def build(self, input_shape: Union[tf.TensorShape, List[tf.TensorShape]]): """Creates the child layers of the layer.""" depthwise_conv2d_quantized = helper.quantize_wrapped_layer( tf.keras.layers.DepthwiseConv2D, configs.Default8BitConvQuantizeConfig( ['depthwise_kernel'], ['activation'], True)) conv2d_quantized = helper.quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig( ['kernel'], [], self._last_quantize)) dwconv_kwargs = self._conv_kwargs.copy() # Depthwise conv input filters is always equal to output filters. # This filters argument only needed for the point-wise conv2d op. del dwconv_kwargs['filters'] self.dw_conv = depthwise_conv2d_quantized(name='dw', **dwconv_kwargs) conv_kwargs = self._conv_kwargs.copy() conv_kwargs.update({ 'kernel_size': (1, 1), 'strides': (1, 1), 'padding': 'valid', 'groups': 1, }) self.conv = conv2d_quantized(name='pw', **conv_kwargs)
def build(self, input_shape): conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], False)) conv2d_quantized_output_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], True)) num_reduced_filters = nn_layers.make_divisible( max(1, int(self._in_filters * self._se_ratio)), divisor=self._divisible_by, round_down_protect=self._round_down_protect) self._se_reduce = conv2d_quantized( filters=num_reduced_filters, kernel_size=1, strides=1, padding='same', use_bias=True, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._se_expand = conv2d_quantized_output_quantized( filters=self._out_filters, kernel_size=1, strides=1, padding='same', use_bias=True, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._multiply = tfmot.quantization.keras.QuantizeWrapperV2( tf.keras.layers.Multiply(), configs.Default8BitQuantizeConfig([], [], True)) self._reduce_mean_quantizer = ( tfmot.quantization.keras.quantizers.MovingAverageQuantizer( num_bits=8, per_axis=False, symmetric=False, narrow_range=False)) self._reduce_mean_quantizer_vars = self._reduce_mean_quantizer.build( None, 'reduce_mean_quantizer_vars', self) self._activation_layer = tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig()) self._create_gating_activation_layer() self._build_quantizer_vars() super().build(input_shape)
def build(self, input_shape: Optional[Union[Sequence[int], tf.Tensor]]): """Build variables and child layers to prepare for calling.""" if self._use_explicit_padding and self._kernel_size > 1: padding_size = nn_layers.get_padding_for_kernel_size( self._kernel_size) self._pad = tf.keras.layers.ZeroPadding2D(padding_size) conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], not self._use_normalization)) self._conv0 = conv2d_quantized( filters=self._filters, kernel_size=self._kernel_size, strides=self._strides, padding=self._padding, use_bias=self._use_bias, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) if self._use_normalization: self._norm0 = self._norm_by_activation(self._activation)( axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon) self._activation_layer = tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig()) super(Conv2DBNBlockQuantized, self).build(input_shape)
return tfmot.quantization.keras.QuantizeWrapperV2( cls(*arg, **kwargs), quantize_config) return constructor def norm_by_activation(activation, norm_quantized, norm_no_quantized): if activation not in ['relu', 'relu6']: return norm_quantized else: return norm_no_quantized Conv2DQuantized = quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], False)) Conv2DOutputQuantized = quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], True)) DepthwiseConv2DQuantized = quantize_wrapped_layer( tf.keras.layers.DepthwiseConv2D, configs.Default8BitConvQuantizeConfig(['depthwise_kernel'], ['activation'], False)) DepthwiseConv2DOutputQuantized = quantize_wrapped_layer( tf.keras.layers.DepthwiseConv2D, configs.Default8BitConvQuantizeConfig(['depthwise_kernel'], ['activation'], True)) GlobalAveragePooling2DQuantized = quantize_wrapped_layer( tf.keras.layers.GlobalAveragePooling2D, configs.Default8BitQuantizeConfig([], [], True)) AveragePooling2DQuantized = quantize_wrapped_layer(
def build(self, input_shape: Union[tf.TensorShape, List[tf.TensorShape]]): """Creates the variables of the head.""" if self._config_dict['use_separable_conv']: conv_op = SeparableConv2DQuantized else: conv_op = helper.quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig( ['kernel'], ['activation'], False)) conv_kwargs = { 'filters': self._config_dict['num_filters'], 'kernel_size': 3, 'padding': 'same', 'bias_initializer': tf.zeros_initializer(), 'bias_regularizer': self._config_dict['bias_regularizer'], } if not self._config_dict['use_separable_conv']: conv_kwargs.update({ 'kernel_initializer': tf.keras.initializers.RandomNormal( stddev=0.01), 'kernel_regularizer': self._config_dict['kernel_regularizer'], }) base_bn_op = (tf.keras.layers.experimental.SyncBatchNormalization if self._config_dict['use_sync_bn'] else tf.keras.layers.BatchNormalization) bn_op = helper.norm_by_activation( self._config_dict['activation'], helper.quantize_wrapped_layer( base_bn_op, configs.Default8BitOutputQuantizeConfig()), helper.quantize_wrapped_layer( base_bn_op, configs.NoOpQuantizeConfig())) bn_kwargs = { 'axis': self._bn_axis, 'momentum': self._config_dict['norm_momentum'], 'epsilon': self._config_dict['norm_epsilon'], } # Class net. self._cls_convs = [] self._cls_norms = [] for level in range( self._config_dict['min_level'], self._config_dict['max_level'] + 1): this_level_cls_norms = [] for i in range(self._config_dict['num_convs']): if level == self._config_dict['min_level']: cls_conv_name = 'classnet-conv_{}'.format(i) self._cls_convs.append(conv_op(name=cls_conv_name, **conv_kwargs)) cls_norm_name = 'classnet-conv-norm_{}_{}'.format(level, i) this_level_cls_norms.append(bn_op(name=cls_norm_name, **bn_kwargs)) self._cls_norms.append(this_level_cls_norms) classifier_kwargs = { 'filters': ( self._config_dict['num_classes'] * self._config_dict['num_anchors_per_location']), 'kernel_size': 3, 'padding': 'same', 'bias_initializer': tf.constant_initializer(-np.log((1 - 0.01) / 0.01)), 'bias_regularizer': self._config_dict['bias_regularizer'], } if not self._config_dict['use_separable_conv']: classifier_kwargs.update({ 'kernel_initializer': tf.keras.initializers.RandomNormal(stddev=1e-5), 'kernel_regularizer': self._config_dict['kernel_regularizer'], }) self._classifier = conv_op( name='scores', last_quantize=True, **classifier_kwargs) # Box net. self._box_convs = [] self._box_norms = [] for level in range( self._config_dict['min_level'], self._config_dict['max_level'] + 1): this_level_box_norms = [] for i in range(self._config_dict['num_convs']): if level == self._config_dict['min_level']: box_conv_name = 'boxnet-conv_{}'.format(i) self._box_convs.append(conv_op(name=box_conv_name, **conv_kwargs)) box_norm_name = 'boxnet-conv-norm_{}_{}'.format(level, i) this_level_box_norms.append(bn_op(name=box_norm_name, **bn_kwargs)) self._box_norms.append(this_level_box_norms) box_regressor_kwargs = { 'filters': (self._config_dict['num_params_per_anchor'] * self._config_dict['num_anchors_per_location']), 'kernel_size': 3, 'padding': 'same', 'bias_initializer': tf.zeros_initializer(), 'bias_regularizer': self._config_dict['bias_regularizer'], } if not self._config_dict['use_separable_conv']: box_regressor_kwargs.update({ 'kernel_initializer': tf.keras.initializers.RandomNormal( stddev=1e-5), 'kernel_regularizer': self._config_dict['kernel_regularizer'], }) self._box_regressor = conv_op( name='boxes', last_quantize=True, **box_regressor_kwargs) # Attribute learning nets. if self._config_dict['attribute_heads']: self._att_predictors = {} self._att_convs = {} self._att_norms = {} for att_config in self._config_dict['attribute_heads']: att_name = att_config['name'] att_type = att_config['type'] att_size = att_config['size'] att_convs_i = [] att_norms_i = [] # Build conv and norm layers. for level in range(self._config_dict['min_level'], self._config_dict['max_level'] + 1): this_level_att_norms = [] for i in range(self._config_dict['num_convs']): if level == self._config_dict['min_level']: att_conv_name = '{}-conv_{}'.format(att_name, i) att_convs_i.append(conv_op(name=att_conv_name, **conv_kwargs)) att_norm_name = '{}-conv-norm_{}_{}'.format(att_name, level, i) this_level_att_norms.append(bn_op(name=att_norm_name, **bn_kwargs)) att_norms_i.append(this_level_att_norms) self._att_convs[att_name] = att_convs_i self._att_norms[att_name] = att_norms_i # Build the final prediction layer. att_predictor_kwargs = { 'filters': (att_size * self._config_dict['num_anchors_per_location']), 'kernel_size': 3, 'padding': 'same', 'bias_initializer': tf.zeros_initializer(), 'bias_regularizer': self._config_dict['bias_regularizer'], } if att_type == 'regression': att_predictor_kwargs.update( {'bias_initializer': tf.zeros_initializer()}) elif att_type == 'classification': att_predictor_kwargs.update({ 'bias_initializer': tf.constant_initializer(-np.log((1 - 0.01) / 0.01)) }) else: raise ValueError( 'Attribute head type {} not supported.'.format(att_type)) if not self._config_dict['use_separable_conv']: att_predictor_kwargs.update({ 'kernel_initializer': tf.keras.initializers.RandomNormal(stddev=1e-5), 'kernel_regularizer': self._config_dict['kernel_regularizer'], }) self._att_predictors[att_name] = conv_op( name='{}_attributes'.format(att_name), **att_predictor_kwargs) super().build(input_shape)
def build(self, input_shape: Optional[Union[Sequence[int], tf.Tensor]]): """Build variables and child layers to prepare for calling.""" conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], False)) depthwise_conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.DepthwiseConv2D, configs.Default8BitConvQuantizeConfig(['depthwise_kernel'], ['activation'], False)) expand_filters = self._in_filters if self._expand_ratio > 1: # First 1x1 conv for channel expansion. expand_filters = nn_layers.make_divisible( self._in_filters * self._expand_ratio, self._divisible_by) expand_kernel = 1 if self._use_depthwise else self._kernel_size expand_stride = 1 if self._use_depthwise else self._strides self._conv0 = conv2d_quantized( filters=expand_filters, kernel_size=expand_kernel, strides=expand_stride, padding='same', use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm0 = self._norm_by_activation(self._activation)( axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon) self._activation_layer = tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig()) if self._use_depthwise: # Depthwise conv. self._conv1 = depthwise_conv2d_quantized( kernel_size=(self._kernel_size, self._kernel_size), strides=self._strides, padding='same', depth_multiplier=1, dilation_rate=self._dilation_rate, use_bias=False, depthwise_initializer=self._kernel_initializer, depthwise_regularizer=self._depthsize_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm1 = self._norm_by_activation(self._depthwise_activation)( axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon) self._depthwise_activation_layer = ( tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._depthwise_activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig())) # Squeeze and excitation. if self._se_ratio and self._se_ratio > 0 and self._se_ratio <= 1: logging.info('Use Squeeze and excitation.') in_filters = self._in_filters if self._expand_se_in_filters: in_filters = expand_filters self._squeeze_excitation = qat_nn_layers.SqueezeExcitationQuantized( in_filters=in_filters, out_filters=expand_filters, se_ratio=self._se_ratio, divisible_by=self._divisible_by, round_down_protect=self._se_round_down_protect, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=self._se_inner_activation, gating_activation=self._se_gating_activation) else: self._squeeze_excitation = None # Last 1x1 conv. self._conv2 = conv2d_quantized( filters=self._out_filters, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm2 = self._norm_with_quantize(axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon) if self._stochastic_depth_drop_rate: self._stochastic_depth = nn_layers.StochasticDepth( self._stochastic_depth_drop_rate) else: self._stochastic_depth = None self._add = tfmot.quantization.keras.QuantizeWrapperV2( tf.keras.layers.Add(), configs.Default8BitQuantizeConfig([], [], True)) super(InvertedBottleneckBlockQuantized, self).build(input_shape)
def build(self, input_shape: Optional[Union[Sequence[int], tf.Tensor]]): """Build variables and child layers to prepare for calling.""" conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], False)) if self._use_projection: if self._resnetd_shortcut: self._shortcut0 = tf.keras.layers.AveragePooling2D( pool_size=2, strides=self._strides, padding='same') self._shortcut1 = conv2d_quantized( filters=self._filters * 4, kernel_size=1, strides=1, use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) else: self._shortcut = conv2d_quantized( filters=self._filters * 4, kernel_size=1, strides=self._strides, use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm0 = self._norm_with_quantize( axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon, trainable=self._bn_trainable) self._conv1 = conv2d_quantized( filters=self._filters, kernel_size=1, strides=1, use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm1 = self._norm(axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon, trainable=self._bn_trainable) self._activation1 = tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig()) self._conv2 = conv2d_quantized( filters=self._filters, kernel_size=3, strides=self._strides, dilation_rate=self._dilation_rate, padding='same', use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm2 = self._norm(axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon, trainable=self._bn_trainable) self._activation2 = tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig()) self._conv3 = conv2d_quantized( filters=self._filters * 4, kernel_size=1, strides=1, use_bias=False, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activation=NoOpActivation()) self._norm3 = self._norm_with_quantize(axis=self._bn_axis, momentum=self._norm_momentum, epsilon=self._norm_epsilon, trainable=self._bn_trainable) self._activation3 = tfmot.quantization.keras.QuantizeWrapperV2( tf_utils.get_activation(self._activation, use_keras_layer=True), configs.Default8BitActivationQuantizeConfig()) if self._se_ratio and self._se_ratio > 0 and self._se_ratio <= 1: self._squeeze_excitation = qat_nn_layers.SqueezeExcitationQuantized( in_filters=self._filters * 4, out_filters=self._filters * 4, se_ratio=self._se_ratio, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer) else: self._squeeze_excitation = None if self._stochastic_depth_drop_rate: self._stochastic_depth = nn_layers.StochasticDepth( self._stochastic_depth_drop_rate) else: self._stochastic_depth = None self._add = tfmot.quantization.keras.QuantizeWrapperV2( tf.keras.layers.Add(), configs.Default8BitQuantizeConfig([], [], True)) super(BottleneckBlockQuantized, self).build(input_shape)
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 = _quantize_wrapped_layer( norm_layer, configs.Default8BitOutputQuantizeConfig()) norm = norm_with_quantize if self._activation not in [ 'relu', 'relu6' ] else _quantize_wrapped_layer(norm_layer, configs.NoOpQuantizeConfig()) conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], False)) depthwise_conv2d_quantized_output_quantized = _quantize_wrapped_layer( tf.keras.layers.DepthwiseConv2D, configs.Default8BitConvQuantizeConfig(['depthwise_kernel'], ['activation'], True)) self.aspp_layers = [] conv1 = conv2d_quantized(filters=self._output_channels, kernel_size=(1, 1), kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_bias=False, activation=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 += [ depthwise_conv2d_quantized_output_quantized( 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=NoOpActivation()) ] kernel_size = (1, 1) conv_dilation = leading_layers + [ conv2d_quantized(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=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 = [ _quantize_wrapped_layer( tf.keras.layers.GlobalAveragePooling2D, configs.Default8BitQuantizeConfig([], [], True))(), _quantize_wrapped_layer( tf.keras.layers.Reshape, configs.Default8BitQuantizeConfig([], [], True))((1, 1, channels)) ] else: pooling = [ _quantize_wrapped_layer( tf.keras.layers.AveragePooling2D, configs.Default8BitQuantizeConfig([], [], True))( self._pool_kernel_size) ] conv2 = conv2d_quantized(filters=self._output_channels, kernel_size=(1, 1), kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_bias=False, activation=NoOpActivation()) norm2 = norm(axis=self._bn_axis, momentum=self._batchnorm_momentum, epsilon=self._batchnorm_epsilon) self.aspp_layers.append(pooling + [conv2, norm2]) resizing = _quantize_wrapped_layer( tf.keras.layers.Resizing, configs.Default8BitQuantizeConfig([], [], True)) self._resizing_layer = resizing(height, width, interpolation=self._interpolation) self._projection = [ conv2d_quantized(filters=self._output_channels, kernel_size=(1, 1), kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_bias=False, activation=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) concat = _quantize_wrapped_layer( tf.keras.layers.Concatenate, configs.Default8BitQuantizeConfig([], [], True)) self._concat_layer = concat(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) conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], False)) conv2d_quantized_output_quantized = _quantize_wrapped_layer( tf.keras.layers.Conv2D, configs.Default8BitConvQuantizeConfig(['kernel'], ['activation'], True)) depthwise_conv2d_quantized = _quantize_wrapped_layer( tf.keras.layers.DepthwiseConv2D, configs.Default8BitConvQuantizeConfig(['depthwise_kernel'], ['activation'], False)) 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 = _quantize_wrapped_layer( norm_layer, configs.Default8BitOutputQuantizeConfig()) norm = norm_with_quantize if self._config_dict['activation'] not in [ 'relu', 'relu6' ] else _quantize_wrapped_layer(norm_layer, configs.NoOpQuantizeConfig()) 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 = conv2d_quantized( 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=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( depthwise_conv2d_quantized( 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=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( conv2d_quantized(name=conv_name, filters=self._config_dict['num_filters'], activation=NoOpActivation(), **conv_kwargs)) norm_name = 'segmentation_head_norm_{}'.format(i) self._norms.append(norm(name=norm_name, **bn_kwargs)) self._classifier = conv2d_quantized_output_quantized( 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=NoOpActivation()) upsampling = _quantize_wrapped_layer( tf.keras.layers.UpSampling2D, configs.Default8BitQuantizeConfig([], [], True)) self._upsampling_layer = upsampling( 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') concat = _quantize_wrapped_layer( tf.keras.layers.Concatenate, configs.Default8BitQuantizeConfig([], [], True)) self._concat_layer = concat(axis=self._bn_axis) super().build(input_shape)