This repository contains a Pytorch implementation of our paper titled Unifying Gradients to Improve Real-world Robustness for Deep Networks.

Using anaconda and install packages with requirements.txt.

$ pip install -r envs/requirements.txt

The first time you run this repository, you need to extract subsets of CIFAR10 and ImageNet using function select_cifar10_100 and select_imagenet_1000, see main.py and datasets.py. A slight modification on the dataloader function is required: annotate subdataloader part if you do not extract it before. The subdataset used in our paper can be found here.

Our module can be easily plugged into any well-trained model. The only required thing is to split the forward process into two parts: get_feature() and get_pred() functions. The model used in our paper can be downloaded here.

Five score-base query attacks, five defense methods can be tested.

python main.py --dataset cifar10 --batch_size 256 --gpu 0 --max_query 100 --p 0.05 --model_type UniG --delta 0.5 --epochs_ 1 --lr_ 10 --eval
python main.py --dataset imagenet --batch_size 32 --gpu 0 --max_query 100 --p 0.8 --model_type UniG --delta 0.1 --epochs_ 1 --lr_ 1 --eval
# model_type：vanilla,AT,RND,DENT,PNI,UniG

Hereby, we choose Square attack, PreResNet18, CIFAR10 as a simple example. The attack setting is: eps=8/255, query=100/2500, p=0.05. More results on different attacks and datasets see paper.

If you have any problem with this code, please contact me directly. Email address: yingwen_wu@sjtu.edu.cn