This is the official repository of the ICML2021 paper Backdoor Scanning for Deep Neural Networks through K-Arm Optimization in PyTorch.
If you find this code is useful for your research, please cite the following:
@article{shen2021backdoor,
title={Backdoor Scanning for Deep Neural Networks through K-Arm Optimization},
author={Shen, Guangyu and Liu, Yingqi and Tao, Guanhong and An, Shengwei and Xu, Qiuling and Cheng, Siyuan and Ma, Shiqing and Zhang, Xiangyu},
journal={arXiv preprint arXiv:2102.05123},
year={2021}
}
We suggest to use Conda for testing the code on TrojAI datasets. Detailed instruction can be found here.
https://www.anaconda.com/distribution/
-
conda create --name trojai-example python=3.8 -y
-
conda activate trojai-example
-
Install required packages into this conda environment
conda install pytorch=1.7.0 torchvision=0.8.0 torchtext==0.8.0 cudatoolkit=11.0 -c pytorch -c conda-forge
pip install --upgrade trojai
conda install jsonpickle
The code is tested on TrojAI datasets (round1-4). TrojAI datasets can be accessed at TrojAI website
We will release the ImageNet pre-trained models and code in the near future.
To test the code on TrojAI datasets, simply run command
$ python main.py --result_filepath <resultFilepath> --examples_dirpath <dataDirpath> --model_filepath <modelFilepath>
To run the code on custom models, make sure your sample images in the <examples_dirpath>
have the following format:
examples_dirpath/class_<class_id>_example_<example_id>.png
Description about the main parameters:
<Beta>
: Coefficient in the K-Arm schedulor objective function<gamma,global(local)_theta>
: Parameters in the Arm Pre-screening procedure<global(local,ratio)_det_bound>
: Trigger size bound for detecting different types of backdoors<epsilon_for_bandits>
: Controls the randomness during the K-Arm optimization
Please check our results here under the team name Perspecta-PurdueRutgers
. The default settings of the parameters in this repo can achieve 90% detection accuracy on TrojAI round3 training and testing datasets. For scanning different types of models, some parameters might need tunning.
Guangyu Shen, shen447@purdue.edu
Yingqi Liu, liu1751@purdue.edu