init

stylegan2/layers/bias_act.py

@@ -0,0 +1,73 @@
+import numpy as np
+import tensorflow as tf
+
+
+# the custom cuda compiled fused_bias_act is not compatible with tf.GradientTape()
+# I think it's because of 2nd order gradient computation
+class BiasAct(tf.keras.layers.Layer):
+    def __init__(self, lrmul, act, **kwargs):
+        super(BiasAct, self).__init__(**kwargs)
+        assert act in ['linear', 'lrelu']
+        self.lrmul = lrmul
+
+        if act == 'linear':
+            self.act = tf.keras.layers.Lambda(lambda x: tf.identity(x))
+            self.gain = 1.0
+        else:
+            self.act = tf.keras.layers.LeakyReLU(alpha=0.2)
+            self.gain = np.sqrt(2)
+
+    def build(self, input_shape):
+        self.len2 = True if len(input_shape) == 2 else False
+        b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32)
+        self.b = tf.Variable(b_init, name='b', trainable=True)
+
+    def call(self, inputs, training=None, mask=None):
+        b = self.lrmul * self.b
+
+        if self.len2:
+            x = inputs + b
+        else:
+            x = inputs + tf.reshape(b, shape=[1, -1, 1, 1])
+        x = self.act(x)
+        x = self.gain * x
+        return x
+
+    def get_config(self):
+        config = super(BiasAct, self).get_config()
+        config.update({
+            'lrmul': self.lrmul,
+            'gain': self.gain,
+            'len2': self.len2,
+        })
+        return config
+
+
+# import tensorflow as tf
+#
+#
+# from stylegan2.layers.cuda.fused_bias_act import fused_bias_act
+#
+#
+# class BiasAct(tf.keras.layers.Layer):
+#     def __init__(self, lrmul, act, **kwargs):
+#         super(BiasAct, self).__init__(**kwargs)
+#         self.lrmul = lrmul
+#         self.act = act
+#
+#     def build(self, input_shape):
+#         b_init = tf.zeros(shape=(input_shape[1],), dtype=tf.dtypes.float32)
+#         self.b = tf.Variable(b_init, name='b', trainable=True)
+#
+#     def call(self, inputs, training=None, mask=None):
+#         b = self.lrmul * self.b
+#         x = fused_bias_act(inputs, b=b, act=self.act, alpha=None, gain=None)
+#         return x
+#
+#     def get_config(self):
+#         config = super(BiasAct, self).get_config()
+#         config.update({
+#             'lrmul': self.lrmul,
+#             'act': self.act,
+#         })
+#         return config