├── .gitignore ├── LICENSE ├── README.md └── nets ├── __init__.py └── dropblock.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | __pycache__/ 3 | 4 | train.py 5 | classification_net.py 6 | dropblock_eager.py 7 | dropblock_test.py 8 | 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 An Jiaoyang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DropBlock in TensorFlow 2 | 3 | This is a TensorFlow implementation of the following paper: 4 | 5 | >DropBlock: A regularization method for convolutional networks 6 | >arXiv. https://arxiv.org/abs/1810.12890 7 | 8 | 9 | ## Usage 10 | 11 | Graph Execution 12 | - For 2D input 13 | ```python 14 | import numpy as np 15 | import tensorflow as tf 16 | from nets.dropblock import DropBlock2D 17 | 18 | # only support `channels_last` data format 19 | a = tf.placeholder(tf.float32, [None, 10, 10, 3]) 20 | keep_prob = tf.placeholder(tf.float32) 21 | training = tf.placeholder(tf.bool) 22 | 23 | drop_block = DropBlock2D(keep_prob=keep_prob, block_size=3) 24 | b = drop_block(a, training) 25 | 26 | sess = tf.Session() 27 | feed_dict = {a: np.ones([2, 10, 10, 3]), keep_prob: 0.8, training: True} 28 | c = sess.run(b, feed_dict=feed_dict) 29 | 30 | print(c[0, :, :, 0]) 31 | ``` 32 | 33 | - For 3D input 34 | ```python 35 | import numpy as np 36 | import tensorflow as tf 37 | from nets.dropblock import DropBlock3D 38 | 39 | # only support `channels_last` data format 40 | a = tf.placeholder(tf.float32, [None, 5, 5, 5, 1]) 41 | keep_prob = tf.placeholder(tf.float32) 42 | training = tf.placeholder(tf.bool) 43 | 44 | drop_block = DropBlock3D(keep_prob=keep_prob, block_size=3) 45 | b = drop_block(a, training) 46 | 47 | sess = tf.Session() 48 | feed_dict = {a: np.ones([1, 5, 5, 5, 1]), keep_prob: 0.2, training: True} 49 | c = sess.run(b, feed_dict=feed_dict) 50 | 51 | for i in range(5): 52 | print(c[0, i, :, :, 0]) 53 | ``` 54 | 55 | Eager Execution 56 | - For 2D input 57 | ```python 58 | import tensorflow as tf 59 | from nets.dropblock import DropBlock2D 60 | 61 | tf.enable_eager_execution() 62 | 63 | # only support `channels_last` data format 64 | a = tf.ones([2, 10, 10, 3]) 65 | 66 | drop_block = DropBlock2D(keep_prob=0.8, block_size=3) 67 | b = drop_block(a, training=True) 68 | 69 | print(b[0, :, :, 0]) 70 | 71 | # update keep probability 72 | drop_block.set_keep_prob(0.1) 73 | b = drop_block(a, training=True) 74 | 75 | print(b[0, :, :, 0]) 76 | ``` 77 | 78 | - For 3D input 79 | ```python 80 | import tensorflow as tf 81 | from nets.dropblock import DropBlock3D 82 | 83 | tf.enable_eager_execution() 84 | 85 | # only support `channels_last` data format 86 | a = tf.ones([[1, 5, 5, 5, 1]]) 87 | 88 | drop_block = DropBlock3D(keep_prob=0.2, block_size=3) 89 | b = drop_block(a, training=True) 90 | 91 | print(b[0, :, :, 0]) 92 | 93 | # update keep probability 94 | drop_block.set_keep_prob(0.1) 95 | b = drop_block(a, training=True) 96 | 97 | print(b[0, :, :, 0]) 98 | ``` 99 | -------------------------------------------------------------------------------- /nets/__init__.py: -------------------------------------------------------------------------------- 1 | # Author: An Jiaoyang 2 | # ============================= 3 | -------------------------------------------------------------------------------- /nets/dropblock.py: -------------------------------------------------------------------------------- 1 | # Author: An Jiaoyang 2 | # ============================= 3 | import tensorflow as tf 4 | from tensorflow.python.keras import backend as K 5 | 6 | 7 | def _bernoulli(shape, mean): 8 | return tf.nn.relu(tf.sign(mean - tf.random_uniform(shape, minval=0, maxval=1, dtype=tf.float32))) 9 | 10 | 11 | class DropBlock2D(tf.keras.layers.Layer): 12 | def __init__(self, keep_prob, block_size, scale=True, **kwargs): 13 | super(DropBlock2D, self).__init__(**kwargs) 14 | self.keep_prob = float(keep_prob) if isinstance(keep_prob, int) else keep_prob 15 | self.block_size = int(block_size) 16 | self.scale = tf.constant(scale, dtype=tf.bool) if isinstance(scale, bool) else scale 17 | 18 | def compute_output_shape(self, input_shape): 19 | return input_shape 20 | 21 | def build(self, input_shape): 22 | assert len(input_shape) == 4 23 | _, self.h, self.w, self.channel = input_shape.as_list() 24 | # pad the mask 25 | p1 = (self.block_size - 1) // 2 26 | p0 = (self.block_size - 1) - p1 27 | self.padding = [[0, 0], [p0, p1], [p0, p1], [0, 0]] 28 | self.set_keep_prob() 29 | super(DropBlock2D, self).build(input_shape) 30 | 31 | def call(self, inputs, training=None, **kwargs): 32 | def drop(): 33 | mask = self._create_mask(tf.shape(inputs)) 34 | output = inputs * mask 35 | output = tf.cond(self.scale, 36 | true_fn=lambda: output * tf.to_float(tf.size(mask)) / tf.reduce_sum(mask), 37 | false_fn=lambda: output) 38 | return output 39 | 40 | if training is None: 41 | training = K.learning_phase() 42 | output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)), 43 | true_fn=lambda: inputs, 44 | false_fn=drop) 45 | return output 46 | 47 | def set_keep_prob(self, keep_prob=None): 48 | """This method only supports Eager Execution""" 49 | if keep_prob is not None: 50 | self.keep_prob = keep_prob 51 | w, h = tf.to_float(self.w), tf.to_float(self.h) 52 | self.gamma = (1. - self.keep_prob) * (w * h) / (self.block_size ** 2) / \ 53 | ((w - self.block_size + 1) * (h - self.block_size + 1)) 54 | 55 | def _create_mask(self, input_shape): 56 | sampling_mask_shape = tf.stack([input_shape[0], 57 | self.h - self.block_size + 1, 58 | self.w - self.block_size + 1, 59 | self.channel]) 60 | mask = _bernoulli(sampling_mask_shape, self.gamma) 61 | mask = tf.pad(mask, self.padding) 62 | mask = tf.nn.max_pool(mask, [1, self.block_size, self.block_size, 1], [1, 1, 1, 1], 'SAME') 63 | mask = 1 - mask 64 | return mask 65 | 66 | 67 | class DropBlock3D(tf.keras.layers.Layer): 68 | def __init__(self, keep_prob, block_size, scale=True, **kwargs): 69 | super(DropBlock3D, self).__init__(**kwargs) 70 | self.keep_prob = float(keep_prob) if isinstance(keep_prob, int) else keep_prob 71 | self.block_size = int(block_size) 72 | self.scale = tf.constant(scale, dtype=tf.bool) if isinstance(scale, bool) else scale 73 | 74 | def compute_output_shape(self, input_shape): 75 | return input_shape 76 | 77 | def build(self, input_shape): 78 | assert len(input_shape) == 5 79 | _, self.d, self.h, self.w, self.channel = input_shape.as_list() 80 | # pad the mask 81 | p1 = (self.block_size - 1) // 2 82 | p0= (self.block_size - 1) - p1 83 | self.padding = [[0, 0], [p0, p1], [p0, p1], [p0, p1], [0, 0]] 84 | self.set_keep_prob() 85 | super(DropBlock3D, self).build(input_shape) 86 | 87 | def call(self, inputs, training=None, **kwargs): 88 | def drop(): 89 | mask = self._create_mask(tf.shape(inputs)) 90 | output = inputs * mask 91 | output = tf.cond(self.scale, 92 | true_fn=lambda: output * tf.to_float(tf.size(mask)) / tf.reduce_sum(mask), 93 | false_fn=lambda: output) 94 | return output 95 | 96 | if training is None: 97 | training = K.learning_phase() 98 | output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)), 99 | true_fn=lambda: inputs, 100 | false_fn=drop) 101 | return output 102 | 103 | def set_keep_prob(self, keep_prob=None): 104 | """This method only supports Eager Execution""" 105 | if keep_prob is not None: 106 | self.keep_prob = keep_prob 107 | d, w, h = tf.to_float(self.d), tf.to_float(self.w), tf.to_float(self.h) 108 | self.gamma = ((1. - self.keep_prob) * (d * w * h) / (self.block_size ** 3) / 109 | ((d - self.block_size + 1) * (w - self.block_size + 1) * (h - self.block_size + 1))) 110 | 111 | def _create_mask(self, input_shape): 112 | sampling_mask_shape = tf.stack([input_shape[0], 113 | self.d - self.block_size + 1, 114 | self.h - self.block_size + 1, 115 | self.w - self.block_size + 1, 116 | self.channel]) 117 | mask = _bernoulli(sampling_mask_shape, self.gamma) 118 | mask = tf.pad(mask, self.padding) 119 | mask = tf.nn.max_pool3d(mask, [1, self.block_size, self.block_size, self.block_size, 1], [1, 1, 1, 1, 1], 'SAME') 120 | mask = 1 - mask 121 | return mask 122 | 123 | --------------------------------------------------------------------------------