简介

PyHessian 是一个基于 PyTorch 的库，用于对神经网络模型进行基于 Hessian 矩阵的分析。该库能够计算以下指标：

• Hessian 矩阵的最大特征值
• Hessian 矩阵的迹
• Hessian 矩阵特征值谱密度（ESD）

更多详细信息请参阅：

• Hessian 教程笔记本
• 教程视频讲解
• PyHessian 论文。

发表论文列表

本项目得到了 NSF 的资助，我们希望记录所有使用或借助 PyHessian 编写的相关论文。这将有助于我们持续改进和开发该库，同时也为二阶方法的相关研究及正在进行的工作提供一份良好的总结。当前的论文列表请见此处。如果您有相关的论文，请与我们联系，我们将非常乐意将其加入列表中。

使用方法

通过 pip 安装

您可以使用 pip 安装该库：

pip install pyhessian

从源码安装

您也可以从源码编译并安装该库：

git clone https://github.com/amirgholami/PyHessian.git
python setup.py install

在运行 Hessian 相关代码之前，我们需要一个（预训练的）神经网络模型。这里我们提供了一个用于在 Cifar-10 数据集上训练 ResNet20 模型的脚本：

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]

可选参数：
--batch-size                训练批次大小（默认：128）
--test-batch-size           测试批次大小（默认：256）
--epochs                    总训练轮数（默认：180）
--lr                        初始学习率（默认：0.1）
--lr-decay                  学习率衰减比例（默认：0.1）
--lr-decay-epoch            学习率衰减的轮次（默认：80、120）
--seed                      用于复现结果的随机种子（默认：1）
--weight-decay              权重衰减系数（默认：5e-4）
--batch-norm                是否在 ResNet 中使用批归一化（默认：是）
--residual                  是否使用残差连接（默认：是）
--cuda                      是否使用 GPU（默认：是）
--saving-folder             最终检查点的保存路径（默认：checkpoints/）

模型检查点保存完成后，我们可以运行以下代码来获取 Hessian 矩阵的最大特征值、迹以及特征值谱密度：

export CUDA_VISIBLE_DEVICES=0; python example_pyhessian_analysis.py [--mini-hessian-batch-size] [--hessian-batch-size] [--seed] [--batch-norm] [--residual] [--cuda] [--resume]

可选参数：
--mini-hessian-batch-size   小批量 Hessian 计算的批次大小（默认：200）
--hessian-batch-size        Hessian 矩阵计算的批次大小（默认：200）
--seed                      用于复现结果的随机种子（默认：1）
--batch-norm                是否在 ResNet 中使用批归一化（默认：是）
--residual                  是否使用残差连接（默认：是）
--cuda                      是否使用 GPU（默认：是）
--resume                    检查点的恢复路径（默认：无，需用户填写）

输出的谱密度图将保存为 example.pdf。

引用

PyHessian 是作为以下论文的一部分而开发的。如果您在工作中发现该库很有用，我们非常感谢您能引用以下论文：

• Z. Yao, A. Gholami, K Keutzer, M. Mahoney. PyHessian: Neural Networks Through the Lens of the Hessian, Spotlight at ICML workshop on Beyond First-Order Optimization Methods in Machine Learning, 2020 (also in proceedings of IEEE Conferenceon big data), PDF。

版权

本软件和/或数据已于 2023 年 2 月 27 日存入 BAIR 开放研究共同体仓库。

PyHessian 快速上手指南

PyHessian 是一个基于 PyTorch 的库，用于对神经网络模型进行黑塞矩阵（Hessian）分析。它支持计算顶部特征值、黑塞矩阵的迹以及完整的特征值谱密度（ESD）。

环境准备

• 操作系统: Linux / macOS (Windows 需自行配置 CUDA 环境)
• 核心依赖:
  • Python 3.6+
  • PyTorch (建议安装与你的 CUDA 版本匹配的最新版)
  • CUDA (可选，但强烈建议使用 GPU 加速计算)
• 前置知识: 需要预先训练好一个 PyTorch 模型（如 ResNet），并保存为 checkpoint 文件。

安装步骤

你可以选择通过 Pip 直接安装，或从源码编译安装。国内用户建议使用清华或阿里镜像源加速下载。

方式一：通过 Pip 安装（推荐）

pip install pyhessian -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

git clone https://github.com/amirgholami/PyHessian.git
cd PyHessian
python setup.py install

基本使用

使用流程分为两步：首先训练模型（或加载已有模型），然后运行分析脚本获取黑塞矩阵指标。

1. 训练模型（示例）

如果你还没有预训练模型，可以使用库中提供的脚本在 CIFAR-10 数据集上训练 ResNet20。以下命令将模型保存在 checkpoints/ 目录：

export CUDA_VISIBLE_DEVICES=0; python training.py --batch-size 128 --epochs 180 --lr 0.1 --weight-decay 5e-4 --cuda True --saving-folder checkpoints/

注意：请根据你的硬件情况调整 --batch-size 或其他超参数。训练完成后，确保 checkpoints/ 目录下生成了模型权重文件。

2. 执行黑塞矩阵分析

模型准备就绪后，运行分析脚本计算顶部特征值、迹和谱密度。请将 --resume 参数替换为你实际的模型权重路径。

export CUDA_VISIBLE_DEVICES=0; python example_pyhessian_analysis.py --mini-hessian-batch-size 200 --hessian-batch-size 200 --cuda True --resume checkpoints/your_model_checkpoint.pth

关键参数说明：

• --resume: 必填，指定预训练模型 checkpoint 的路径。
• --hessian-batch-size: 用于计算黑塞矩阵的批次大小（默认 200，显存不足时可调小）。
• --cuda: 是否使用 GPU 加速（默认 True）。

运行结束后，当前目录下将生成名为 example.pdf 的文件，其中包含特征值谱密度的可视化图表。

提示：更多高级用法、API 详解及理论背景，请参考项目自带的 Hessian_Tutorial.ipynb 教程笔记。