bnelearn.util.autograd_hacks module

Library for extracting interesting quantites from autograd, see README.md

Not thread-safe because of module-level variables

Author: cybertronai @ https://github.com/cybertronai/autograd-hacks.

Notation: o: number of output classes (exact Hessian), number of Hessian samples (sampled Hessian) n: batch-size do: output dimension (output channels for convolution) di: input dimension (input channels for convolution) Hi: per-example Hessian of matmul, shaped as matrix of [dim, dim], indices have been row-vectorized Hi_bias: per-example Hessian of bias Oh, Ow: output height, output width (convolution) Kh, Kw: kernel height, kernel width (convolution)

Jb: batch output Jacobian of matmul, output sensitivity for example,class pair, [o, n, ….] Jb_bias: as above, but for bias

A, activations: inputs into current layer B, backprops: backprop values (aka Lop aka Jacobian-vector product) observed at current layer

bnelearn.util.autograd_hacks.add_hooks(model: Module) → None

Adds hooks to model to save activations and backprop values.

The hooks will 1. save activations into param.activations during forward pass 2. append backprops to params.backprops_list during backward pass.

Call “remove_hooks(model)” to disable this.

Args:

model:

bnelearn.util.autograd_hacks.backprop_hess(output: Tensor, hess_type: str) → None

Call backprop 1 or more times to get values needed for Hessian computation.

Args:

output: prediction of neural network (ie, input of nn.CrossEntropyLoss()) hess_type: type of Hessian propagation, “CrossEntropy” results in exact Hessian for CrossEntropy

Returns:

bnelearn.util.autograd_hacks.clear_backprops(model: Module) → None

Delete layer.backprops_list in every layer.

bnelearn.util.autograd_hacks.compute_grad1(model: Module, loss_type: str = 'mean') → None

Compute per-example gradients and save them under ‘param.grad1’. Must be called after loss.backprop()

Args:

model: loss_type: either “mean” or “sum” depending whether backpropped loss was averaged or summed over batch

bnelearn.util.autograd_hacks.compute_hess(model: Module) → None

Save Hessian under param.hess for each param in the model

bnelearn.util.autograd_hacks.disable_hooks() → None

Globally disable all hooks installed by this library.

bnelearn.util.autograd_hacks.enable_hooks() → None

the opposite of disable_hooks()

bnelearn.util.autograd_hacks.is_supported(layer: Module) → bool

Check if this layer is supported

bnelearn.util.autograd_hacks.remove_hooks(model: Module) → None

Remove hooks added by add_hooks(model)

bnelearn.util.autograd_hacks.symsqrt(a, cond=None, return_rank=False, dtype=torch.float32)

Symmetric square root of a positive semi-definite matrix. See https://github.com/pytorch/pytorch/issues/25481