Merge pull request #49 from Scyfer/added_dropout_switch

ethereon · web-flow · commit 0c5bf53817f8 · 2016-08-07T22:00:51.000-07:00
Added a switch variable which allows it to disable dropout
diff --git a/kaffe/tensorflow/network.py b/kaffe/tensorflow/network.py
@@ -40,6 +40,10 @@ def __init__(self, inputs, trainable=True):
         self.layers = dict(inputs)
         # If true, the resulting variables are set as trainable
         self.trainable = trainable
+        # Switch variable for dropout
+        self.use_dropout = tf.placeholder_with_default(tf.constant(1.0),
+                                                       shape=[],
+                                                       name='use_dropout')
         self.setup()
 
     def setup(self):
@@ -236,4 +240,5 @@ def batch_normalization(self, input, name, scale_offset=True, relu=False):
 
     @layer
     def dropout(self, input, keep_prob, name):
-        return tf.nn.dropout(input, keep_prob, name=name)
+        keep = 1 - self.use_dropout + (self.use_dropout * keep_prob)
+        return tf.nn.dropout(input, keep, name=name)
