-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdropblock.py
124 lines (103 loc) · 5.3 KB
/
dropblock.py
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
# Author: An Jiaoyang
# =============================
import tensorflow as tf
from tensorflow.python.keras import backend as K
def _bernoulli(shape, mean):
return tf.nn.relu(tf.sign(mean - tf.random_uniform(shape, minval=0, maxval=1, dtype=tf.float32)))
class DropBlock2D(tf.keras.layers.Layer):
def __init__(self, keep_prob, block_size, scale=True, **kwargs):
super(DropBlock2D, self).__init__(**kwargs)
self.keep_prob = float(keep_prob) if isinstance(keep_prob, int) else keep_prob
self.block_size = int(block_size)
self.scale = tf.constant(scale, dtype=tf.bool) if isinstance(scale, bool) else scale
def compute_output_shape(self, input_shape):
return input_shape
def build(self, input_shape):
assert len(input_shape) == 4
_, self.h, self.w, self.channel = input_shape.as_list()
# pad the mask
p1 = (self.block_size - 1) // 2
p0 = (self.block_size - 1) - p1
self.padding = [[0, 0], [p0, p1], [p0, p1], [0, 0]]
self.set_keep_prob()
super(DropBlock2D, self).build(input_shape)
def call(self, inputs, training=None, **kwargs):
def drop():
mask = self._create_mask(tf.shape(inputs))
output = inputs * mask
output = tf.cond(self.scale,
true_fn=lambda: output * tf.to_float(tf.size(mask)) / tf.reduce_sum(mask),
false_fn=lambda: output)
return output
if training is None:
training = K.learning_phase()
output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)),
true_fn=lambda: inputs,
false_fn=drop)
return output
def set_keep_prob(self, keep_prob=None):
"""This method only supports Eager Execution"""
if keep_prob is not None:
self.keep_prob = keep_prob
w, h = tf.to_float(self.w), tf.to_float(self.h)
self.gamma = (1. - self.keep_prob) * (w * h) / (self.block_size ** 2) / \
((w - self.block_size + 1) * (h - self.block_size + 1))
def _create_mask(self, input_shape):
sampling_mask_shape = tf.stack([input_shape[0],
self.h - self.block_size + 1,
self.w - self.block_size + 1,
self.channel])
mask = _bernoulli(sampling_mask_shape, self.gamma)
mask = tf.pad(mask, self.padding)
mask = tf.nn.max_pool(mask, [1, self.block_size, self.block_size, 1], [1, 1, 1, 1], 'SAME')
mask = 1 - mask
return mask
class DropBlock3D(tf.keras.layers.Layer):
def __init__(self, keep_prob, block_size, scale=True, **kwargs):
super(DropBlock3D, self).__init__(**kwargs)
self.keep_prob = float(keep_prob) if isinstance(keep_prob, int) else keep_prob
self.block_size = int(block_size)
self.scale = tf.constant(scale, dtype=tf.bool) if isinstance(scale, bool) else scale
def compute_output_shape(self, input_shape):
return input_shape
def build(self, input_shape):
assert len(input_shape) == 5
_, self.d, self.h, self.w, self.channel = input_shape.as_list()
# pad the mask
p1 = (self.block_size - 1) // 2
p0= (self.block_size - 1) - p1
self.padding = [[0, 0], [p0, p1], [p0, p1], [p0, p1], [0, 0]]
self.set_keep_prob()
super(DropBlock3D, self).build(input_shape)
def call(self, inputs, training=None, **kwargs):
def drop():
mask = self._create_mask(tf.shape(inputs))
output = inputs * mask
output = tf.cond(self.scale,
true_fn=lambda: output * tf.to_float(tf.size(mask)) / tf.reduce_sum(mask),
false_fn=lambda: output)
return output
if training is None:
training = K.learning_phase()
output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)),
true_fn=lambda: inputs,
false_fn=drop)
return output
def set_keep_prob(self, keep_prob=None):
"""This method only supports Eager Execution"""
if keep_prob is not None:
self.keep_prob = keep_prob
d, w, h = tf.to_float(self.d), tf.to_float(self.w), tf.to_float(self.h)
self.gamma = ((1. - self.keep_prob) * (d * w * h) / (self.block_size ** 3) /
((d - self.block_size + 1) * (w - self.block_size + 1) * (h - self.block_size + 1)))
def _create_mask(self, input_shape):
sampling_mask_shape = tf.stack([input_shape[0],
self.d - self.block_size + 1,
self.h - self.block_size + 1,
self.w - self.block_size + 1,
self.channel])
mask = _bernoulli(sampling_mask_shape, self.gamma)
mask = tf.pad(mask, self.padding)
mask = tf.nn.max_pool3d(mask, [1, self.block_size, self.block_size, self.block_size, 1], [1, 1, 1, 1, 1], 'SAME')
mask = 1 - mask
return mask