File size: 4,990 Bytes
506da10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
# 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.
"""Drop path operation.
This scripts implements the drop path operation, proposed in
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, Kilian Weinberger,
Deep Networks with Stochastic Depth. In ECCV, 2016.
"""
import tensorflow as tf
def get_drop_path_keep_prob(keep_prob_for_last_stage, schedule,
current_stage, num_stages):
"""Gets drop path keep probability for current stage.
Args:
keep_prob_for_last_stage: A float, the drop path keep probability for
last stage. This flag is used in conjunction with the flag `schedule`, as
they together determine drop path keep probability for the other stages.
schedule: A string, the drop path schedule. Currently, we support
'constant': use the same drop path keep probability for all stages, and
'linear': linearly decrease the drop path keep probability from 1.0 at
0-th stage (or STEM) to `keep_prob_for_last_stage` at last stage.
current_stage: An integer, current stage number.
num_stages: An integer, the number of stages.
Returns:
The drop path keep probability for the current stage.
Raises:
ValueError: If schedule is not supported.
"""
if schedule == 'constant':
return keep_prob_for_last_stage
elif schedule == 'linear':
return 1.0 - (1.0 - keep_prob_for_last_stage) * current_stage / num_stages
else:
raise ValueError('Unexpected schedule %s.' % schedule)
def generate_drop_path_random_mask(input_tensor, drop_path_keep_prob):
"""Generates a random mask for drop path.
This function generates a random mask for training models with drop path. Each
scalar in the output indicates whether the block or path will be kept. The
scalars are scaled with (1.0 / drop_path_keep_prob) so that the output will
have the same expectation no mather what the drop_path_keep_prob is.
Reference:
"Deep Networks with Stochastic Depth" https://arxiv.org/pdf/1603.09382.pdf
Args:
input_tensor: An input [batch_size, n_1, n_2, ..., n_k] tensor.
drop_path_keep_prob: A float, the keep probability for dropping path.
Returns:
binary_tensor: A [batch_size, 1, 1, ..., 1] tensor with the same dtype as
the input_tensor.
"""
binary_tensor = None
if drop_path_keep_prob < 1.0:
input_shape = input_tensor.get_shape().as_list()
random_tensor_shape = [input_shape[0]] + [1] * (len(input_shape) - 1)
random_tensor = drop_path_keep_prob
random_tensor += tf.random.uniform(
random_tensor_shape, dtype=input_tensor.dtype)
binary_tensor = tf.math.divide(tf.floor(random_tensor), drop_path_keep_prob)
return binary_tensor
class DropPath(tf.keras.layers.Layer):
"""Drop path layer.
For details, please see the original paper listed below.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, Kilian Weinberger,
Deep Networks with Stochastic Depth. In ECCV, 2016.
"""
def __init__(self, drop_path_keep_prob=1.0, name=None):
"""Initializes a drop path layer.
Args:
drop_path_keep_prob: A float, the keep probability for dropping path.
name: An optional string specifying the operation name.
Rasies:
ValueError: If drop_path_keep_prob is <= 0 or > 1.
"""
super(DropPath, self).__init__(name=name)
self._drop_path_keep_prob = drop_path_keep_prob
if self._drop_path_keep_prob <= 0 or self._drop_path_keep_prob > 1.0:
raise ValueError('drop_path_keep_prob not valid. Got %f.' %
self._drop_path_keep_prob)
def call(self, input_tensor, training=False):
"""Performs a forward pass.
Args:
input_tensor: An input tensor of type tf.Tensor with shape [batch, height,
width, channels].
training: A boolean flag indicating whether training behavior should be
used (default: False).
Returns:
The output tensor.
"""
if self._drop_path_keep_prob == 1.0 or not training:
return input_tensor
drop_path_random_mask = generate_drop_path_random_mask(
input_tensor, self._drop_path_keep_prob)
if drop_path_random_mask is not None:
input_tensor = input_tensor * drop_path_random_mask
return input_tensor
def get_config(self):
config = {
'drop_path_keep_prob': self._drop_path_keep_prob,
}
base_config = super(DropPath, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
|