Adversarial Neuron Pruning Purifies Backdoored Deep Models
Code for NeurIPS 2021 "Adversarial Neuron Pruning Purifies Backdoored Deep Models" by Dongxian Wu and Yisen Wang.
News
11/08/2021 - Our checkpoints and recipe have been released.
10/31/2021 - Our code has be released.
10/28/2021 - Our paper and slide have be released.
10/26/2021 - Our code and paper will be released soon.
What ANP Does
ANP can easily repair backdoored deep models using limited clean data and limited computational resources. Only 500 clean images from CIFAR-10 and 2000 iterations are used in the displayed example.
Requisite
This code is implemented in PyTorch, and we have tested the code under the following environment settings:
- python = 3.7.3
- torch = 1.8.0
- torchvision = 0.9.0
A Quick Start - How to use it
For a detailed introduction, please refer to our recipe.
Step 1: Train a backdoored DNN
By default, we train a backdoored resnet-18 under badnets with 5% poison rate and class 0 as target label,
python train_backdoor_cifar.py --output-dir './save'
We save trained backdoored model and the trigger info as ./save/last_model.th
and ./save/trigger_info.th
. Some checkpoints have been released in Google drive or Baidu drive (pwd: bmrb).
Step 2: Optimize masks under neuron perturbations
We optimize the mask for each neuron under neuron perturbations, and save mask values in './save/mask_values.txt' . By default, we only use 500 clean data to optimize.
python optimize_mask_cifar.py --output-dir './save' --checkpoints './save/last_model.th' --trigger-info' './save/trigger_info.th'
Step 3: Prune neurons to defend
You can prune neurons by threshold,
python prune_neuron_cifar.py --output-dir './save' --mask-file './save/mask_values.txt' --checkpoints './save/last_model.th' --trigger-info' './save/trigger_info.th'
Citing this work
If you use our code, please consider cite the following: Dongxian Wu and Yisen Wang. Adversarial Neuron Pruning Purifies Backdoored Deep Models. In NeurIPS, 2021.
@inproceedings{wu2021adversarial,
title={Adversarial Neuron Pruning Purifies Backdoored Deep Models},
author={Dongxian Wu and Yisen Wang},
booktitle={NeurIPS},
year={2021}
}
If there is any problem, be free to open an issue or contact: [email protected].
Useful Links
[1] Mode Connectivity Repair (MCR) defense: https://github.com/IBM/model-sanitization/tree/master/backdoor
[2] Input-aware Backdoor (IAB) attack: https://github.com/VinAIResearch/input-aware-backdoor-attack-release