deeplab.py 8.08 KB
Newer Older
Yeqing Li's avatar
Yeqing Li committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Zhenyu Tan's avatar
Zhenyu Tan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Yeqing Li's avatar
Yeqing Li committed
14

Zhenyu Tan's avatar
Zhenyu Tan committed
15
16
17
18
19
"""Layers for DeepLabV3."""

import tensorflow as tf


20
class SpatialPyramidPooling(tf.keras.layers.Layer):
Zhenyu Tan's avatar
Zhenyu Tan committed
21
22
  """Implements the Atrous Spatial Pyramid Pooling.

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
23
  References:
Zhenyu Tan's avatar
Zhenyu Tan committed
24
25
    [Rethinking Atrous Convolution for Semantic Image Segmentation](
      https://arxiv.org/pdf/1706.05587.pdf)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
26
27
    [Encoder-Decoder with Atrous Separable Convolution for Semantic Image
    Segmentation](https://arxiv.org/pdf/1802.02611.pdf)
Zhenyu Tan's avatar
Zhenyu Tan committed
28
29
30
31
32
33
  """

  def __init__(
      self,
      output_channels,
      dilation_rates,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
34
      pool_kernel_size=None,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
35
      use_sync_bn=False,
Zhenyu Tan's avatar
Zhenyu Tan committed
36
      batchnorm_momentum=0.99,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
37
      batchnorm_epsilon=0.001,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
38
      activation='relu',
Zhenyu Tan's avatar
Zhenyu Tan committed
39
40
41
42
      dropout=0.5,
      kernel_initializer='glorot_uniform',
      kernel_regularizer=None,
      interpolation='bilinear',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
43
      use_depthwise_convolution=False,
Zhenyu Tan's avatar
Zhenyu Tan committed
44
      **kwargs):
45
    """Initializes `SpatialPyramidPooling`.
Zhenyu Tan's avatar
Zhenyu Tan committed
46

47
    Args:
48
      output_channels: Number of channels produced by SpatialPyramidPooling.
Zhenyu Tan's avatar
Zhenyu Tan committed
49
      dilation_rates: A list of integers for parallel dilated conv.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
50
51
52
      pool_kernel_size: A list of integers or None. If None, global average
        pooling is applied, otherwise an average pooling of pool_kernel_size
        is applied.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
53
      use_sync_bn: A bool, whether or not to use sync batch normalization.
Zhenyu Tan's avatar
Zhenyu Tan committed
54
55
      batchnorm_momentum: A float for the momentum in BatchNorm. Defaults to
        0.99.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
56
57
      batchnorm_epsilon: A float for the epsilon value in BatchNorm. Defaults to
        0.001.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
58
      activation: A `str` for type of activation to be used. Defaults to 'relu'.
Zhenyu Tan's avatar
Zhenyu Tan committed
59
60
61
62
63
64
      dropout: A float for the dropout rate before output. Defaults to 0.5.
      kernel_initializer: Kernel initializer for conv layers. Defaults to
        `glorot_uniform`.
      kernel_regularizer: Kernel regularizer for conv layers. Defaults to None.
      interpolation: The interpolation method for upsampling. Defaults to
        `bilinear`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
65
66
67
68
      use_depthwise_convolution: Allows spatial pooling to be separable
         depthwise convolusions. [Encoder-Decoder with Atrous Separable
         Convolution for Semantic Image Segmentation](
         https://arxiv.org/pdf/1802.02611.pdf)
Zhenyu Tan's avatar
Zhenyu Tan committed
69
70
      **kwargs: Other keyword arguments for the layer.
    """
71
    super(SpatialPyramidPooling, self).__init__(**kwargs)
Zhenyu Tan's avatar
Zhenyu Tan committed
72
73
74

    self.output_channels = output_channels
    self.dilation_rates = dilation_rates
Abdullah Rashwan's avatar
Abdullah Rashwan committed
75
    self.use_sync_bn = use_sync_bn
Zhenyu Tan's avatar
Zhenyu Tan committed
76
    self.batchnorm_momentum = batchnorm_momentum
Abdullah Rashwan's avatar
Abdullah Rashwan committed
77
    self.batchnorm_epsilon = batchnorm_epsilon
Abdullah Rashwan's avatar
Abdullah Rashwan committed
78
    self.activation = activation
Zhenyu Tan's avatar
Zhenyu Tan committed
79
80
81
82
83
    self.dropout = dropout
    self.kernel_initializer = tf.keras.initializers.get(kernel_initializer)
    self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)
    self.interpolation = interpolation
    self.input_spec = tf.keras.layers.InputSpec(ndim=4)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
84
    self.pool_kernel_size = pool_kernel_size
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
85
    self.use_depthwise_convolution = use_depthwise_convolution
Zhenyu Tan's avatar
Zhenyu Tan committed
86
87
88
89
90
91
92
93

  def build(self, input_shape):
    height = input_shape[1]
    width = input_shape[2]
    channels = input_shape[3]

    self.aspp_layers = []

Abdullah Rashwan's avatar
Abdullah Rashwan committed
94
95
96
97
98
99
100
101
102
103
    if self.use_sync_bn:
      bn_op = tf.keras.layers.experimental.SyncBatchNormalization
    else:
      bn_op = tf.keras.layers.BatchNormalization

    if tf.keras.backend.image_data_format() == 'channels_last':
      bn_axis = -1
    else:
      bn_axis = 1

Zhenyu Tan's avatar
Zhenyu Tan committed
104
105
106
107
    conv_sequential = tf.keras.Sequential([
        tf.keras.layers.Conv2D(
            filters=self.output_channels, kernel_size=(1, 1),
            kernel_initializer=self.kernel_initializer,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
108
109
110
111
112
113
            kernel_regularizer=self.kernel_regularizer,
            use_bias=False),
        bn_op(
            axis=bn_axis,
            momentum=self.batchnorm_momentum,
            epsilon=self.batchnorm_epsilon),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
114
        tf.keras.layers.Activation(self.activation)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
115
    ])
Zhenyu Tan's avatar
Zhenyu Tan committed
116
117
118
    self.aspp_layers.append(conv_sequential)

    for dilation_rate in self.dilation_rates:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
119
120
121
122
123
124
125
126
127
128
129
130
      leading_layers = []
      kernel_size = (3, 3)
      if self.use_depthwise_convolution:
        leading_layers += [
            tf.keras.layers.DepthwiseConv2D(
                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)
        ]
        kernel_size = (1, 1)
      conv_sequential = tf.keras.Sequential(leading_layers + [
Zhenyu Tan's avatar
Zhenyu Tan committed
131
          tf.keras.layers.Conv2D(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
132
              filters=self.output_channels, kernel_size=kernel_size,
Zhenyu Tan's avatar
Zhenyu Tan committed
133
134
135
              padding='same', kernel_regularizer=self.kernel_regularizer,
              kernel_initializer=self.kernel_initializer,
              dilation_rate=dilation_rate, use_bias=False),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
136
137
          bn_op(axis=bn_axis, momentum=self.batchnorm_momentum,
                epsilon=self.batchnorm_epsilon),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
138
          tf.keras.layers.Activation(self.activation)])
Zhenyu Tan's avatar
Zhenyu Tan committed
139
140
      self.aspp_layers.append(conv_sequential)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
    if self.pool_kernel_size is None:
      pool_sequential = tf.keras.Sequential([
          tf.keras.layers.GlobalAveragePooling2D(),
          tf.keras.layers.Reshape((1, 1, channels))])
    else:
      pool_sequential = tf.keras.Sequential([
          tf.keras.layers.AveragePooling2D(self.pool_kernel_size)])

    pool_sequential.add(
        tf.keras.Sequential([
            tf.keras.layers.Conv2D(
                filters=self.output_channels,
                kernel_size=(1, 1),
                kernel_initializer=self.kernel_initializer,
                kernel_regularizer=self.kernel_regularizer,
                use_bias=False),
            bn_op(
                axis=bn_axis,
                momentum=self.batchnorm_momentum,
                epsilon=self.batchnorm_epsilon),
            tf.keras.layers.Activation(self.activation),
            tf.keras.layers.experimental.preprocessing.Resizing(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
163
164
165
166
                height,
                width,
                interpolation=self.interpolation,
                dtype=tf.float32)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
167
168
        ]))

Zhenyu Tan's avatar
Zhenyu Tan committed
169
170
171
172
173
174
    self.aspp_layers.append(pool_sequential)

    self.projection = tf.keras.Sequential([
        tf.keras.layers.Conv2D(
            filters=self.output_channels, kernel_size=(1, 1),
            kernel_initializer=self.kernel_initializer,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
175
176
177
178
179
180
            kernel_regularizer=self.kernel_regularizer,
            use_bias=False),
        bn_op(
            axis=bn_axis,
            momentum=self.batchnorm_momentum,
            epsilon=self.batchnorm_epsilon),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
181
        tf.keras.layers.Activation(self.activation),
Zhenyu Tan's avatar
Zhenyu Tan committed
182
183
184
185
186
187
188
        tf.keras.layers.Dropout(rate=self.dropout)])

  def call(self, inputs, training=None):
    if training is None:
      training = tf.keras.backend.learning_phase()
    result = []
    for layer in self.aspp_layers:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
189
      result.append(tf.cast(layer(inputs, training=training), inputs.dtype))
Zhenyu Tan's avatar
Zhenyu Tan committed
190
191
192
193
194
195
196
197
    result = tf.concat(result, axis=-1)
    result = self.projection(result, training=training)
    return result

  def get_config(self):
    config = {
        'output_channels': self.output_channels,
        'dilation_rates': self.dilation_rates,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
198
        'pool_kernel_size': self.pool_kernel_size,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
199
        'use_sync_bn': self.use_sync_bn,
Zhenyu Tan's avatar
Zhenyu Tan committed
200
        'batchnorm_momentum': self.batchnorm_momentum,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
201
        'batchnorm_epsilon': self.batchnorm_epsilon,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
202
        'activation': self.activation,
Zhenyu Tan's avatar
Zhenyu Tan committed
203
204
205
206
207
208
209
        'dropout': self.dropout,
        'kernel_initializer': tf.keras.initializers.serialize(
            self.kernel_initializer),
        'kernel_regularizer': tf.keras.regularizers.serialize(
            self.kernel_regularizer),
        'interpolation': self.interpolation,
    }
210
    base_config = super(SpatialPyramidPooling, self).get_config()
Zhenyu Tan's avatar
Zhenyu Tan committed
211
    return dict(list(base_config.items()) + list(config.items()))