TrojanZoo: Evaluating Neural Backdoor Attacks/Defenses

This is the code implementation (pytorch) for our paper: TROJANZOO

TrojanZoo provides a universal pytorch platform to conduct security researches (especially backdoor attacks/defenses) of image classification in deep learning. It is composed of two packages: trojanzoo and trojanvision. trojanzoo contains abstract classes and utilities, while trojanvision contains abstract and concrete ones for image classification task.

Screenshot

Features

1. Colorful and verbose output!

Note: enable with --color for color and --verbose for verbose.
To open an interactive window with color, use python - --color

2. Modular design (plug and play)
3. Good code linting support (this package requires python>=3.9)
4. Register your own module to the library.
5. Native Pytorch Output

trojanzoo and trojanvision provides API to generate raw pytorch instances, which makes it flexible to work with native pytorch and other 3rd party libraries.

6. Enable pytorch native AMP(Automatic Mixed Precision) with --amp for training
7. Flexible Configuration Files
8. Good help information to check arguments. (-h or --help)
9. Detailed and well-organized summary() for each module.

Installation

pip install trojanzoo

(todo) conda install trojanzoo

docker pull local0state/trojanzoo

Quick Start

You can use the provided example scripts to reproduce the evaluation results in our paper.

Note: The program won’t save results without --save

1. Train a model: e.g. ResNetComp18 on CIFAR10 with 95% Acc

python train.py --verbose 1 --amp --dataset cifar10 --model resnetcomp18 --epoch 300 --lr 0.1 --lr_scheduler --lr_step_size 100 --save

2. Test backdoor attack (e.g., BadNet): e.g. BadNet with ResNetComp18 on CIFAR10

python backdoor_attack.py --verbose 1 --pretrain --validate_interval 1 --amp --dataset cifar10 --model resnetcomp18 --attack badnet --random_init --epoch 50 --lr 0.01 --save

3. Test backdoor defense (e.g., Neural Cleanse): e.g. Neural Cleanse against BadNet

python backdoor_attack.py --verbose 1 --pretrain --validate_interval 1 --dataset cifar10 --model resnetcomp18 --attack badnet --defense neural_cleanse --random_init --epoch 50 --lr 0.01

IMC

python backdoor_attack.py --verbose --pretrain --validate_interval 1 --amp --dataset cifar10 --model resnetcomp18 --attack imc --random_init --epoch 50 --lr 0.01 --save

AdvMind

(with attack adaptive and model adaptive)

python adv_defense.py --verbose --pretrain --validate_interval 1 --dataset cifar10 --model resnetcomp18 --attack pgd --defense advmind --attack_adapt --defense_adapt

Detailed Usage

Configuration file structure

All arguments in the parser are able to set default values in configuration files.
If argument values are not set in the config files, we will use the default values of __init__()

Parameters Config: (priority ascend order)

The higher priority config will override lower priority ones.
Within each priority channel, trojanvision configs will overwrite trojanzoo

1. Package Default: /trojanzoo/configs/, /trojanvision/configs/These are package default settings. Please don’t modify them. You can use this as a template to set other configs.
2. User Default: Not decided yet.(Enable it in the code trojanzoo/configs/__init__.py, trojanvision/configs/__init__.py)
3. Workspace Default: /configs/trojanzoo/, /configs/trojanvision/
4. Custom Config: --config [config location]
5. CMD parameters: --[parameter] [value]

Store path of Dataset, Model, Attack & Defense Results

Modify them in corresponding config files and command-line arguments.

Dataset: --data_dir (./data/data)
Model: --model_dir (./data/model)
Attack: --attack_dir (./data/attack)
Defense: --defense_dir (./data/defense)

Output Verbose Information:

1. CMD modules: --verbose
2. Colorful output: --color
3. tqdm progress bar: --tqdm
4. Check command-line argument usage: --help
5. AdvMind verbose information: --output [number]

Use your DIY Dataset/Model/Attack/Defense

1. Follow our example to write your DIY class. (CIFAR10, ResNet, IMC, Neural Cleanse)It’s necessary to subclass our base class. (Dataset, Model, Attack, Defense)
   Optional base classes depending on your use case: (ImageSet, ImageFolder, ImageModel)
2. Register your DIY class in trojanvisionExample: trojanvision.attacks.class_dict[attack_name]=AttackClass
3. Create your config files if necessary.
   No need to modify any codes. Just directly add {attack_name}.yml (.json) in the config directory.
4. Good to go!

Todo List

1. Sphinx Docs
2. Unit test