# coding=utf-8 # Copyright 2021 The Deeplab2 Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Squeeze and excite layer. This script implements the squeeze-and-excite (SE), proposed in - Squeeze-and-Excitation Networks, Jie Hu, Li Shen, Samuel Albanie, Gang Sun, Enhua Wu. In CVPR 2018. Recently, this SE operation is further simplied with a single fully connected layer, referred as simplified_squeeze_and_excite in our implementation. For details, please see - Lee and Park proposed to use only one fully connected layer in SE. CenterMask : Real-Time Anchor-Free Instance Segmentation. Youngwan Lee and Jongyoul Park. In CVPR 2020. """ from typing import Optional from absl import logging import tensorflow as tf from deeplab2.model import utils from deeplab2.model.layers import activations layers = tf.keras.layers class SimplifiedSqueezeAndExcite(tf.keras.layers.Layer): """A simplified squeeze-and-excite layer. Original squeeze-and-exciation (SE) is proposed in Squeeze-and-Excitation Networks, Jie Hu, Li Shen, Samuel Albanie, Gang Sun, Enhua Wu. In CVPR 2018. Lee and Park proposed to use only one fully connected layer in SE. CenterMask : Real-Time Anchor-Free Instance Segmentation. Youngwan Lee and Jongyoul Park. In CVPR 2020. In this function, we implement the simplified version of SE. Additionally, we follow MobileNetv3 to use the hard sigmoid function. """ def __init__(self, squeeze_channels, name=None): """Initializes a simplified squeeze-and-excite layer. Args: squeeze_channels: Integer, channels for the squeezed features. name: An optional string specifying the operation name. """ super(SimplifiedSqueezeAndExcite, self).__init__(name=name) self._squeeze_channels = squeeze_channels self._se_conv = layers.Conv2D(self._squeeze_channels, 1, name='squeeze_and_excite', use_bias=True, kernel_initializer='VarianceScaling') self._hard_sigmoid = activations.get_activation('hard_sigmoid') def call(self, input_tensor): """Performs a forward pass. Args: input_tensor: An input tensor of type tf.Tensor with shape [batch, height, width, channels]. Returns: The output tensor. """ pooled = tf.reduce_mean(input_tensor, [1, 2], keepdims=True) squeezed = self._se_conv(pooled) excited = self._hard_sigmoid(squeezed) * input_tensor return excited def get_config(self): config = { 'squeeze_channels': self._squeeze_channels, } base_config = super(SimplifiedSqueezeAndExcite, self).get_config() return dict(list(base_config.items()) + list(config.items())) class SqueezeAndExcite(tf.keras.layers.Layer): """Creates a squeeze and excitation layer. Reference: Squeeze-and-Excitation Networks, Jie Hu, Li Shen, Samuel Albanie, Gang Sun, Enhua Wu. In CVPR 2018. This implementation follows the original SE and differs from the above simplified version. """ def __init__( self, in_filters: int, out_filters: int, se_ratio: float, divisible_by: int = 1, kernel_initializer: str = 'VarianceScaling', kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None, bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None, activation: str = 'relu', gating_activation: str = 'sigmoid', name: Optional[str] = None): """Initializes a squeeze and excitation layer. Args: in_filters: The number of filters that se_ratio should be applied to. out_filters: The number of filters of the output tensor. se_ratio: The SE ratio for the squeeze and excitation layer. divisible_by: An `int` that ensures all inner dimensions are divisible by this number. kernel_initializer: The kernel_initializer for convolutional layers. kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D. Default to None. bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2d. Default to None. activation: The name of the activation function. gating_activation: The name of the activation function for final gating function. name: The layer name. """ super(SqueezeAndExcite, self).__init__(name=name) self._in_filters = in_filters self._out_filters = out_filters self._se_ratio = se_ratio self._divisible_by = divisible_by self._activation = activation self._gating_activation = gating_activation self._kernel_initializer = kernel_initializer self._kernel_regularizer = kernel_regularizer self._bias_regularizer = bias_regularizer if tf.keras.backend.image_data_format() == 'channels_last': self._spatial_axis = [1, 2] else: self._spatial_axis = [2, 3] self._activation_fn = activations.get_activation(activation) self._gating_activation_fn = activations.get_activation(gating_activation) num_reduced_filters = utils.make_divisible( max(1, int(self._in_filters * self._se_ratio)), divisor=self._divisible_by) if self._se_ratio > 1.0: logging.warn('Squeezing ratio %d is larger than 1.0.', self._se_ratio) self._se_reduce = tf.keras.layers.Conv2D( 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, name=name + '_reduce') self._se_expand = tf.keras.layers.Conv2D( 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, name=name + '_expand') def call(self, inputs): x = tf.reduce_mean(inputs, self._spatial_axis, keepdims=True) x = self._activation_fn(self._se_reduce(x)) x = self._gating_activation_fn(self._se_expand(x)) return x * inputs