Spaces:
Runtime error
Runtime error
| # 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 | |