init

stylegan2/layers/embedding.py

@@ -0,0 +1,24 @@
+import tensorflow as tf
+
+
+class LabelEmbedding(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, **kwargs):
+        super(LabelEmbedding, self).__init__(**kwargs)
+        self.embed_dim = embed_dim
+
+    def build(self, input_shape):
+        weight_shape = [input_shape[1], self.embed_dim]
+        # tf 1.15 mean(0.0), std(1.0) default value of tf.initializers.random_normal()
+        w_init = tf.random.normal(shape=weight_shape, mean=0.0, stddev=1.0)
+        self.w = tf.Variable(w_init, name='w', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        x = tf.matmul(inputs, self.w)
+        return x
+
+    def get_config(self):
+        config = super(LabelEmbedding, self).get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+        })
+        return config