26 lines
799 B
Python
26 lines
799 B
Python
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
|