Awesome
Efficiently Computing Local Lipschitz Constants of Neural Networks via Bound Propagation
Code for paper:
Efficiently Computing Local Lipschitz Constants of Neural Networks via Bound Propagation, by Zhouxing Shi, Yihan Wang, Huan Zhang, Zico Kolter and Cho-Jui Hsieh. In NeurIPS 2022.
The core implementation of this paper is now a part of auto_LiRPA. See the example about bounding Jacobian, Jacobian-vector product, and Linf local Lipschitz constants.
For reproducing our results, please install auto_LiRPA version 0.3.1 which was released in November, 2022. We are working on a more general and flexible support for Jacobian in auto_LiRPA.
Dependencies
Python 3.7+ and PyTorch 1.11+ are recommended.
Install other Python libraries:
pip install -r requirements.txt
Train Models
We first need to train models which will be saved to models_pretrained/
for analyzing local Lipschitz constants:
python train_model.py --data DATA --model MODEL
DATA
can be chosen from simple
, MNIST
, CIFAR
, and tinyimagenet
.
For tinyimagenet
, data need to be downloaded first with:
cd data/tinyImageNet
bash tinyimagenet_download.sh
MODEL
can be chosen from models available under models/
.
Some models in models/simple.py
have arguments width
and depth
for experiments
with varying width or depth, and they can be set by --width WIDTH
or --depth DEPTH
respectively.
Other options include --num-epochs
and --lr
for setting number of epochs and learning rates.
Compute Local Lipschitz Constants
To compute local Lipschitz constants by our method, we run:
python main.py --data DATA --model MODEL --load PATH_TO_MODEL_FILE --eps EPS
PATH_TO_MODEL_FILE
is the path to the checkpoint of the pretrained model, and EPS
is the radius of input domain. For models from models/simple.py
, their width and depth
can be specified by --model-params width=WIDTH
or --model-params depth=DEPTH
.
By default, BaB is not used. To enable Branch-and-Bound (BaB), add --bab
, and the time
budget can be set by --timeout TIMEOUT
. Batch size of BaB can be set by --batch-size BATCH-SIZE
to fit it into the GPU memory.