variable dropout

1 files changed, 6 insertions, 5 deletions

diff --git a/dcgan.py b/dcgan.py
index 347f61e..4ee5917 100644
--- a/dcgan.py
+++ b/dcgan.py
@@ -17,7 +17,7 @@ import sys
 import numpy as np
 
 class DCGAN():
-    def __init__(self, conv_layers = 1):
+    def __init__(self, conv_layers = 1, dropout=0.25):
         # Input shape
         self.img_rows = 28
         self.img_cols = 28
@@ -25,6 +25,7 @@ class DCGAN():
         self.img_shape = (self.img_rows, self.img_cols, self.channels)
         self.latent_dim = 100
         self.conv_layers = conv_layers
+        self.dropout = dropout
 
         optimizer = Adam(0.002, 0.5)
 
@@ -88,20 +89,20 @@ class DCGAN():
 
         model.add(Conv2D(32, kernel_size=3, strides=2, input_shape=self.img_shape, padding="same"))
         model.add(LeakyReLU(alpha=0.2))
-        model.add(Dropout(0.25))
+        model.add(Dropout(self.dropout))
         model.add(Conv2D(64, kernel_size=3, strides=2, padding="same"))
         model.add(ZeroPadding2D(padding=((0,1),(0,1))))
         model.add(BatchNormalization())
         model.add(LeakyReLU(alpha=0.2))
-        model.add(Dropout(0.25))
+        model.add(Dropout(dropout))
         model.add(Conv2D(128, kernel_size=3, strides=2, padding="same"))
         model.add(BatchNormalization())
         model.add(LeakyReLU(alpha=0.2))
-        model.add(Dropout(0.25))
+        model.add(Dropout(dropout))
         model.add(Conv2D(256, kernel_size=3, strides=1, padding="same"))
         model.add(BatchNormalization())
         model.add(LeakyReLU(alpha=0.2))
-        model.add(Dropout(0.25))
+        model.add(Dropout(dropout))
         model.add(Flatten())
         model.add(Dense(1, activation='sigmoid'))