fix: mode of padding

kozistr · Oct 12, 2020 · d6113fb · d6113fb
1 parent 1d222f3
commit d6113fb
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/awesome_gans/wgan/model.py b/awesome_gans/wgan/model.py
@@ -62,11 +62,11 @@ def __init__(self, config):
     def build_discriminator(self) -> tf.keras.Model:
         inputs = Input((self.width, self.height, self.n_channels))
 
-        x = Conv2D(self.n_feats, 5, 2)(inputs)
+        x = Conv2D(self.n_feats, kernel_size=5, strides=2, padding='same')(inputs)
         x = LeakyReLU(alpha=0.2)(x)
 
         for i in range(3):
-            x = Conv2D(self.n_feats * (2 ** (i + 1)), 5, 2)(x)
+            x = Conv2D(self.n_feats * (2 ** (i + 1)), kernel_size=5, strides=2, padding='same')(x)
             x = BatchNormalization()(x)
             x = LeakyReLU(alpha=0.2)(x)
 
@@ -86,11 +86,11 @@ def build_generator(self) -> tf.keras.Model:
         x = Reshape((4, 4, 4 * self.z_dims))(x)
 
         for i in range(3):
-            x = Conv2DTranspose(self.z_dims * 4 // (2 ** i), kernel_size=5, strides=2)(x)
+            x = Conv2DTranspose(self.z_dims * 4 // (2 ** i), kernel_size=5, strides=2, padding='same')(x)
             x = BatchNormalization()(x)
             x = ReLU()(x)
 
-        x = Conv2DTranspose(self.n_channels, kernel_size=5, strides=1)(x)
+        x = Conv2DTranspose(self.n_channels, kernel_size=5, strides=1, padding='same')(x)
         x = Layer('tanh')(x)
 
         return Model(inputs, x, name='generator')