repository·master·Indexed 21 days ago
https://github.com/amirgholami/pyhessianA PyTorch library for Hessian-based analysis of neural network models. It provides tools to compute eigenvalues, the Hessian trace, and Eigenvalue Spectral Density (ESD) using randomized linear algebra, such as Hutchinson's method and the Stochastic Lanczos algorithm, enabling analysis of loss landscape flatness without explicitly forming the full Hessian matrix.
PyHessian requires a pre-trained neural network model. You can use the provided training.py script to train a ResNet20 model on the Cifar-10 dataset.
Note: Ensure you set CUDA_VISIBLE_DEVICES if using a GPU.
export CUDA_VISIBLE_DEVICES=0; python training.py [--batch-size] [--test-batch-size] [--epochs] [--lr] [--lr-decay] [--lr-decay-epoch] [--seed] [--weight-decay] [--batch-norm] [--residual] [--cuda] [--saving-folder]To install from the source repository, clone the repository and run the setup script.
git clone https://github.com/amirgholami/PyHessian.git
python setup.py installYou can install the PyHessian library directly from PyPI using pip.
pip install pyhessianOnce a model checkpoint is saved, use example_pyhessian_analysis.py to compute the top eigenvalue, the trace of the Hessian, and the Eigenvalue Spectral Density (ESD). The output density plot is saved as example.pdf.
export CUDA_VISIBLE_DEVICES=0; python example_pyhessian_analysis.py [--mini-hessian-batch-size] [--hessian-batch-size] [--seed] [--batch-norm] [--residual] [--cuda] [--resume]You can visualize the loss landscape by perturbing model parameters along specific directions. Using the top Hessian eigenvector provides a more informative view of the landscape's curvature than random directions.
Workflow:
hessian_comp.eigenvalues().Note: When perturbing, ensure you use a copy of the model to avoid modifying the original weights permanently.
# 1. Get top eigenvector
top_eigenvalues, top_eigenvector = hessian_comp.eigenvalues()
# 2. Define perturbation range
lams = np.linspace(-0.5, 0.5, 21).astype(np.float32)
# 3. Perturb and measure loss
loss_list = []
model_perb = copy_of_model
for lam in lams:
# Perturb model_perb along top_eigenvector[0] by amount 'lam'
model_perb = get_params(model, model_perb, top_eigenvector[0], lam)
loss_list.append(criterion(model_perb(inputs), targets).item())Arguments for the training.py script:
| Argument | Description |
|---|---|
--batch-size | training batch size (default: 128) |
--test-batch-size | testing batch size (default: 256) |
--epochs | total number of training epochs (default: 180) |
--lr | initial learning rate (default: 0.1) |
--lr-decay | learning rate decay ratio (default: 0.1) |
--lr-decay-epoch | epoch for the learning rate decaying (default: 80, 120) |
--seed | used to reproduce the results (default: 1) |
--weight-decay | weight decay value (default: 5e-4) |
--batch-norm | do we need batch norm in ResNet or not (default: True) |
--residual | do we need residual connection or not (default: True) |
--cuda | do we use gpu or not (default: True) |
--saving-folder | saving path of the final checkpoint (default: checkpoints/) |
Arguments for the example_pyhessian_analysis.py script:
| Argument | Description |
|---|---|
--mini-hessian-batch-size | mini hessian batch size (default: 200) |
--hessian-batch-size | hessian batch size (default: 200) |
--seed | used to reproduce the results (default: 1) |
--batch-norm | do we need batch norm in ResNet or not (default: True) |
--residual | do we need residual connection or not (default: True) |
--cuda | do we use gpu or not (default: True) |
--resume | resume path of the checkpoint (default: none, must be filled by user) |
The hessian module provides methods to approximate the Hessian's trace and its eigenvalue spectrum density using randomized linear algebra (e.g., Hutchinson's method and Stochastic Lanczos algorithm).
trace(): Returns an approximation of the Hessian trace.density(): Returns the eigenvalue spectrum density, useful for analyzing the flatness of the loss landscape.# Assuming hessian_comp is an instance of the hessian module
# Get the trace
trace = hessian_comp.trace()
print(f"Trace: {np.mean(trace)}")
# Get eigenvalue spectrum density
density_eigen, density_weight = hessian_comp.density()When perturbing model parameters along a direction (like a random vector or a gradient vector), use pyhessian.utils.normalization to ensure the direction vector is properly scaled.
from pyhessian.utils import normalization
# Example: Normalizing a list of parameter tensors
v = [torch.randn_like(p) for p in model.parameters()]
v = normalization(v)Use the hessian function to create a Hessian computation module for a PyTorch model. This module allows you to compute eigenvalues and eigenvectors without explicitly forming the full Hessian matrix, making it efficient for large neural networks.
To use it:
.eval() mode to disable running statistics updates.torch.nn.CrossEntropyLoss()).(inputs, targets).cuda=True for faster computation if a GPU is available.from pyhessian import hessian
# Setup
model.eval()
criterion = torch.nn.CrossEntropyLoss()
# inputs, targets should be tensors
# Create the module
hessian_comp = hessian(model, criterion, data=(inputs, targets), cuda=True)
# Compute top N eigenvalues and eigenvectors
top_eigenvalues, top_eigenvector = hessian_comp.eigenvalues(top_n=2)