OBSan is an implementation of the paper "OBsan: An Out-Of-Bound Sanitizer to Harden DNN Executables," based on TVM.
This repo contains the source code for OBSan, together with the dataset preparation tools, evaluation scripts, and downstream applications used in the paper.
This project requires a working TVM installation. A recommended setup is to
clone the TVM repo into a directory tvm/
and
then clone this repo into a subdirectory tvm/obsan/
. In this setup, the code
organization is as follows:
tvm/ # TVM root
|- obsan/ # OBSan root
|- apps/ # Downstream applications
|- eval/ # Evaluation tools
|- support/ # Support libraries / dataset utilities
|- backward.py
|- cgutils.py
|- ...
First clone the TVM repo and checkout the commit OBSan is developed on:
git clone --recursive https://github.com/apache/tvm tvm
cd tvm
git checkout b6b0bafde
git submodule update
Then build the TVM shared libraries following the official
instructions.
The build output should be under tvm/built/
.
Clone this repo into tvm/obsan/
:
git clone https://github.com/yanzuochen/obsan
The project is developed on Python 3.8.12 and Ubuntu 18.04. To ensure the best compatibility, the setup tools will invoke pyenv to install the same version and initialize a virtual environment with all the dependencies. Run the following commands to automate this step:
cd obsan
./setup.sh
source obsan-env.sh
In this step, we need to download CIFAR10 as the training and validation datasets and ChestX-ray8 as the undefined images dataset. We also run a few scripts to generate the AE and perception-broken datasets for each of the three models (paths in the scripts may need to be changed first):
python ./support/aegen/aegen.py --model <resnet50|googlenet|densenet121>
python ./support/broken.py <resnet50|googlenet|densenet121> <output_dir>
Evaluation scripts can be found in the directory tvm/obsan/eval/
and can be
used for reproducing tables III-VIII as well as tables XI and XII in
appendices. Results will be saved to tvm/obsan/results/
. A description of
each of the evaluation scripts is as follows:
python ./eval/evaluation_base.py # Tables III-V, XIV, XV
python ./eval/evaluation_sel.py # Table VI, XVI
python ./eval/evaluation_bob.py # Tables VII, VIII, XVII
The two downstream applications, namely online AE generation prevention (Sec.
IX.A; Table IX) and feedback-driven fuzzing (Sec. IX.B; Table X), are available
at tvm/obsan/apps/bae/
and tvm/obsan/apps/fuzz.py
, respectively.
To launch the online AE attack in the default scenario which allows a perturbation budget of eps = 0.3 and 50 queries per seed, use the following commands:
./apps/bae/attack.sh none 0.3 50 # Without OBSan
./apps/bae/attack.sh NBC 0.3 50 # With FOBSan
./apps/bae/attack.sh gn2 0.3 50 # With BOBSan
./apps/bae/attack.sh NBC+gn2 0.3 50 # With HOBSan
For the sophisticated scenario with a perturbation budget of 0.035 and 500 queries per seed, use the following commands instead:
./apps/bae/attack.sh none 0.035 500 # Without OBSan
./apps/bae/attack.sh NBC 0.035 500 # With FOBSan
./apps/bae/attack.sh gn2 0.035 500 # With BOBSan
./apps/bae/attack.sh NBC+gn2 0.035 500 # With HOBSan
To launch the fuzzing task and reproduce table X, use the following commands:
./fuzz.py --model <resnet50|googlenet|densenet121> --blind # Blackbox
./fuzz.py --model <resnet50|googlenet|densenet121> # Greybox
Results will be saved to tvm/obsan/results/fuzz/
.
This part of results are already generated by ./eval/evaluation_base.py
in
Step 4.