74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
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
|