add ae sub folder

stylegan2_ae/layers/noise.py

@@ -0,0 +1,25 @@
+import tensorflow as tf
+
+
+class Noise(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super(Noise, self).__init__(**kwargs)
+
+    def build(self, input_shape):
+        w_init = tf.zeros(shape=(), dtype=tf.dtypes.float32)
+        self.noise_strength = tf.Variable(w_init, trainable=True, name='w')
+
+    def call(self, inputs, noise=None, training=None, mask=None):
+        x_shape = tf.shape(inputs)
+
+        # noise: [1, 1, x_shape[2], x_shape[3]] or None
+        if noise is None:
+            noise = tf.random.normal(shape=(x_shape[0], 1, x_shape[2], x_shape[3]), dtype=tf.dtypes.float32)
+
+        x = inputs + noise * self.noise_strength
+        return x
+
+    def get_config(self):
+        config = super(Noise, self).get_config()
+        config.update({})
+        return config