def multiply(self, vector):
left_factor = self._input_factor.get_cov()
right_factor = self._output_factor.get_cov()
reshaped_vector = utils.layer_params_to_mat2d(vector)
reshaped_out = (math_ops.matmul(reshaped_vector, right_factor) +
self._output_damping * reshaped_vector)
reshaped_out = (math_ops.matmul(left_factor, reshaped_out) +
self._input_damping * reshaped_out)
if self._renorm_coeff != 1.0:
reshaped_out *= math_ops.cast(self._renorm_coeff,
dtype=reshaped_out.dtype)
return utils.mat2d_to_layer_params(vector, reshaped_out)
def multiply(self, vector):
"""Approximate damped Fisher-vector product.
Args:
vector: Tensor or 2-tuple of Tensors. if self._has_bias, Tensor of shape
[input_size, output_size] corresponding to layer's weights. If not, a
2-tuple of the former and a Tensor of shape [output_size] corresponding
to the layer's bias.
Returns:
Tensor of the same shape, corresponding to the Fisher-vector product.
"""
reshaped_vect = utils.layer_params_to_mat2d(vector)
reshaped_out = reshaped_vect * (self._factor.get_cov() + self._damping)
return utils.mat2d_to_layer_params(vector, reshaped_out)
def multiply_inverse(self, vector):
# import tensorflow as tf
# with tf.control_dependencies([tf.check_numerics(self._input_damping, message='_input_damping')]):
left_factor_inv = self._input_factor.get_damped_inverse(
self._input_damping)
right_factor_inv = self._output_factor.get_damped_inverse(
self._output_damping)
reshaped_vector = utils.layer_params_to_mat2d(vector)
# with tf.control_dependencies([tf.check_numerics(left_factor_inv, message='left_factor_inv'),
# tf.check_numerics(right_factor_inv, message='right_factor_inv'),
# tf.check_numerics(reshaped_vector, message='reshaped_vector')]):
reshaped_out = math_ops.matmul(
left_factor_inv, math_ops.matmul(reshaped_vector,
right_factor_inv))
if self._renorm_coeff != 1.0:
reshaped_out /= math_ops.cast(self._renorm_coeff,
dtype=reshaped_out.dtype)
return utils.mat2d_to_layer_params(vector, reshaped_out)
def multiply(self, vector):
left_basis = self._input_factor.get_eigen_basis(self._input_damping)
right_basis = self._output_factor.get_eigen_basis(self._output_damping)
scaling = self._scale_factor.get_cov()
reshaped_vector = utils.layer_params_to_mat2d(vector)
# Project gradients in EB.
reshaped_out = \
utils.eigen_basis_kron_product_2d(left_basis, right_basis, reshaped_vector, transpose=True)
# Element-wise scaling.
damped_scale_factor = scaling + self._damping
reshaped_out = reshaped_out * damped_scale_factor
# Project back in parameter basis.
reshaped_out = \
utils.eigen_basis_kron_product_2d(left_basis, right_basis, reshaped_out, transpose=False)
return utils.mat2d_to_layer_params(vector, reshaped_out)
def multiply(self, vector):
reshaped_vect = utils.layer_params_to_mat2d(vector)
reshaped_out = reshaped_vect * (self._factor.get_cov() + self._damping)
return utils.mat2d_to_layer_params(vector, reshaped_out)