def call(self, inputs, training=None):
"""Call override to add noise."""
return core_ops.dense(inputs,
self._noise(self.kernel, training),
self._noise(self.bias, training),
self.activation,
dtype=self._compute_dtype_object)
def func(self, inputs, **kwargs):
x = core_ops.dense(inputs,
self.kernel,
self.bias,
self.activation,
dtype=self._compute_dtype_object)
return tf.nn.dropout(x, noise_shape=None, rate=rate)
def _call_block(self, inputs: tf.Tensor, state: tf.Tensor, name: str, num_of_params: int) -> tf.Tensor:
x = tf.concat([inputs, tf.reshape(state, [self.batch_size, -1])], axis=1)
for i in range(num_of_params):
cur_name_base = f'{name}_{i}'
kernel = vars(self)[f"{cur_name_base}_kernel"]
bias = vars(self)[f"{cur_name_base}_bias"]
x = dense(x, kernel, bias, tf.keras.activations.relu)
return x
def call(self, inputs):
kernel = self.kernel + tf.stop_gradient(
tf.math.multiply(self.kernel, self.kernel_dv) - self.kernel)
bias = self.bias + tf.stop_gradient(
tf.math.multiply(self.bias, self.bias_dv) - self.bias)
return core_ops.dense(inputs,
kernel,
bias,
self.activation,
dtype=self._compute_dtype_object)
def call(self, x):
x = dense(inputs=x,
kernel=self.kernel * self.scale,
bias=self.bias,
activation=self.activation,
dtype=self._compute_dtype_object)
if self.use_bias:
x = tf.nn.bias_add(x, self.bias)
return x
def _dense_call(self, og_call, layer, inputs, *args, **kwargs):
# TODO: Support layers with no bias.
kernel = layer.kernel
bias = layer.bias
# NOTE: We need to put this in two if statements as autograph is dumb.
if kernel.ref() not in self.ref_to_logvar:
# This can happen for non-mergeable variables, so we skip them.
return og_call(layer, inputs, *args, **kwargs)
if not self.sampling:
return og_call(layer, inputs, *args, **kwargs)
kernel_logvar = self.ref_to_logvar[kernel.ref()]
bias_logvar = self.ref_to_logvar[bias.ref()]
output = core_ops.dense(
inputs,
kernel,
bias,
activations.get(None),
dtype=layer._compute_dtype_object,
)
output_var = core_ops.dense(
tf.square(inputs) + 1e-2,
tf.exp(kernel_logvar),
tf.exp(bias_logvar),
activations.get(None),
dtype=layer._compute_dtype_object,
)
eps = tf.random.normal(tf.shape(output))
output += tf.sqrt(output_var) * eps
if layer.activation is not None:
output = layer.activation(output)
return output
def call(self, inputs, **kwargs):
if self.scaling:
bjorckscaling = self.get_safe_bjorck_scaling()
bjorckscaling = bjorckscaling * tf.math.reduce_max(
tf.abs(self.kernel), axis=None)
else:
bjorckscaling = 1.0
ortho_w = bjorck_orthonormalize(self.kernel / bjorckscaling,
beta=self.bjorck_beta,
iters=self.bjorck_iter,
order=self.bjorck_order)
# ortho_w_t = tf.transpose(ortho_w)
return core_ops.dense(inputs,
ortho_w,
self.bias,
self.activation,
dtype=self._compute_dtype_object)
def noised():
kernel_sigma = tf.tensordot(self.sigma[self.p:],
self.sigma[:self.p],
axes=0)
kernel = self.kernel + kernel_sigma * K.random_normal(
shape=(self.q, self.p), mean=0., stddev=1., dtype=inputs.dtype)
bias = self.bias
if self.bias is not None:
bias = bias + self.sigma[:self.p] * K.random_normal(
shape=(self.p, ), mean=0., stddev=1., dtype=inputs.dtype)
if tf.config.list_physical_devices('GPU'):
return core_ops_dense(inputs,
kernel,
bias,
self.activation,
dtype=self._compute_dtype,
units=self.p)
else:
return core_ops.dense(inputs,
kernel,
bias,
self.activation,
dtype=self._compute_dtype_object)
def call(self, inputs):
return core_ops.dense(inputs,
self.kernel,
self.bias,
self.activation,
dtype=self._compute_dtype_object)
