def __inv_from_eigen(self,
eigvals,
eigvecs,
truncate=NUMERICAL_STABILITY_CONSTANT):
inv_eigvals = 1. / eigvals
inv_eigvals.clamp_(min=0., max=1. / truncate)
return einsum('ij,j,kj->ik', (eigvecs, inv_eigvals, eigvecs))
def weight(self, ext, module, g_inp, g_out, backproped):
X, dE_dY = convUtils.get_weight_gradient_factors(
module.input0, g_out[0], module)
if module.has_bias():
X = module.append_ones(X)
return einsum('bml,bkl,bmi,bki->b', (dE_dY, X, dE_dY, X))
def bias(self, ext, module, g_inp, g_out, backproped):
sqrt_h_outs = backproped["matrices"]
sqrt_h_outs_signs = backproped["signs"]
h_diag = torch.zeros_like(module.bias)
for h_sqrt, sign in zip(sqrt_h_outs, sqrt_h_outs_signs):
h_diag.add_(sign * einsum('bic->i', (h_sqrt**2, )))
return h_diag
def extract_weight_diagonal(module, input, grad_output):
"""
input must be the unfolded input to the convolution (see unfold_func)
and grad_output the backpropagated gradient
"""
grad_output_viewed = separate_channels_and_pixels(module, grad_output)
AX = einsum('bkl,bmlc->cbkm', (input, grad_output_viewed))
return (AX**2).sum([0, 1]).transpose(0, 1)
def weight(self, ext, module, g_inp, g_out, backproped):
X, dE_dY = convUtils.get_weight_gradient_factors(
module.input0, g_out[0], module)
if module.has_bias():
X = module.append_ones(X)
d1 = einsum('bml,bkl->bmk', (dE_dY, X))
return (d1**2).sum(0).view_as(module.weight)
def weight(self, ext, module, g_inp, g_out, backproped):
sqrt_h_outs = backproped["matrices"]
sqrt_h_outs_signs = backproped["signs"]
h_diag = torch.zeros_like(module.weight)
for h_sqrt, sign in zip(sqrt_h_outs, sqrt_h_outs_signs):
h_diag.add_(sign * einsum('bic,bj->ij',
(h_sqrt**2, module.input0**2)))
return h_diag
def _factors_from_input(self, ext, module):
X = module.homogeneous_unfolded_input()
batch = X.size(0)
ea_strategy = ext.get_ea_strategy()
if ExpectationApproximation.should_average_param_jac(ea_strategy):
raise NotImplementedError
else:
yield einsum('bik,bjk->ij', (X, X)) / batch
def _factors_from_input(self, ext, module):
X = convUtils.unfold_func(module)(module.input0)
batch = X.size(0)
ea_strategy = ext.get_ea_strategy()
if ExpectationApproximation.should_average_param_jac(ea_strategy):
raise NotImplementedError("Undefined")
else:
yield einsum('bik,bjk->ij', (X, X)) / batch
def bias(self, ext, module, g_inp, g_out, backproped):
sqrt_h_outs = backproped["matrices"]
sqrt_h_outs_signs = backproped["signs"]
h_diag = torch.zeros_like(module.bias)
for h_sqrt, sign in zip(sqrt_h_outs, sqrt_h_outs_signs):
h_sqrt_view = convUtils.separate_channels_and_pixels(
module, h_sqrt)
h_diag.add_(sign * einsum('bijc,bikc->i',
(h_sqrt_view, h_sqrt_view)))
return h_diag
def matrix_from_kron_facs(factors):
assert all_tensors_of_order(order=2, tensors=factors)
mat = None
for factor in factors:
if mat is None:
mat = factor
else:
new_shape = (mat.shape[0] * factor.shape[0],
mat.shape[1] * factor.shape[1])
mat = einsum('ij,kl->ikjl',
(mat, factor)).contiguous().view(new_shape)
return mat
def CMP_in(mat):
"""Multiplication of curvature matrix with matrix `mat`.
Parameters:
-----------
mat : torch.Tensor
Matrix that will be multiplied.
"""
Jmat = self.derivatives.jac_mat_prod(module, g_inp, g_out, mat)
CJmat = CMP_out(Jmat)
JTCJmat = self.derivatives.jac_t_mat_prod(module, g_inp, g_out,
CJmat)
if residual_mod is not None:
JTCJmat.add_(einsum('bi,bic->bic', (residual_mod, mat)))
return JTCJmat
def _factor_from_sqrt(self, backproped):
return [einsum('bic,bjc->ij', (backproped, backproped))]
def __mean_input_outer(self, module):
N, flat_input = self.derivatives.batch_flat(module.input0)
return einsum('bi,bj->ij', (flat_input, flat_input)) / N
def vp(v):
assert len(v.shape) == 1
v_reshaped = v.view(col_dims)
return einsum(equation, v_reshaped, *factors).view(-1)
def weight(self, ext, module, g_inp, g_out, backproped):
return einsum('bi,bj->ij', (g_out[0]**2, module.input0**2))
def _factor_from_sqrt(self, module, backproped):
sqrt_ggn = backproped
sqrt_ggn = convUtils.separate_channels_and_pixels(module, sqrt_ggn)
sqrt_ggn = einsum('bijc->bic', (sqrt_ggn, ))
return einsum('bic,blc->il', (sqrt_ggn, sqrt_ggn))
def weight(self, ext, module, g_inp, g_out, backproped):
input = module.homogeneous_input()
return einsum('bi,bj->ij', (g_out[0]**2, input**2))