pytorch-hessian-eigenthings

by noahgolmant

Efficient PyTorch Hessian eigendecomposition tools!

216 Stars 15 Forks Last release: Not found MIT License 27 Commits 0 Releases

Available items

No Items, yet!

The developer of this repository has not created any items for sale yet. Need a bug fixed? Help with integration? A different license? Create a request here:

pytorch-hessian-eigenthings

The

hessian-eigenthings
module provides an efficient (and scalable!) way to compute the eigendecomposition of the Hessian for an arbitrary PyTorch model. It uses PyTorch's Hessian-vector product and your choice of (a) the Lanczos method or (b) stochastic power iteration with deflation in order to compute the top eigenvalues and eigenvectors of the Hessian.

Installation

For now, you have to install from this repo. It's a tiny thing so why put it on pypi.

pip install --upgrade git+https://github.com/noahgolmant/[email protected]#egg=hessian-eigenthings

Usage

The main function you're probably interested in is

compute_hessian_eigenthings
. Sample usage is like so:
import torch
from hessian_eigenthings import compute_hessian_eigenthings

model = ResNet18() dataloader = ... loss = torch.nn.functional.cross_entropy

num_eigenthings = 20 # compute top 20 eigenvalues/eigenvectors

eigenvals, eigenvecs = compute_hessian_eigenthings(model, dataloader, loss, num_eigenthings)

This also includes a more general power iteration with deflation implementation in

power_iter.py
.
lanczos.py
calls a
scipy
hook
to a battle-tested ARPACK implementation.

Citing this work

If you find this repo useful and would like to cite it in a publication, here is a BibTeX entry:

@misc{hessian-eigenthings,
    author       = {Noah Golmant, Zhewei Yao, Amir Gholami, Michael Mahoney, Joseph Gonzalez},
    title        = {pytorch-hessian-eigentings: efficient PyTorch Hessian eigendecomposition},
    month        = oct,
    year         = 2018,
    version      = {1.0},
    url          = {https://github.com/noahgolmant/pytorch-hessian-eigenthings}
    }

Acknowledgements

This code was written in collaboration with Zhewei Yao, Amir Gholami, Michael Mahoney, and Joseph Gonzalez in UC Berkeley's RISELab.

The deflated power iteration routine is based on code in the HessianFlow repository recently described in the following paper: Z. Yao, A. Gholami, Q. Lei, K. Keutzer, M. Mahoney. "Hessian-based Analysis of Large Batch Training and Robustness to Adversaries", NIPS'18 (arXiv:1802.08241)

Stochastic power iteration with acceleration is based on the following paper: C. De Sa, B. He, I. Mitliagkas, C. Ré, P. Xu. "Accelerated Stochastic Power Iteration", PMLR-21 (arXiv:1707.02670)

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.