a universal pytorch platform to conduct security researches
Project description
TrojanZoo
NOTE: TrojanZoo requires
python>=3.9.2
,pytorch>=1.8.1
andtorchvision>=0.9.1
, which must be installed manually. Recommend to useconda
to install.
This is the code implementation (pytorch) for our paper:
TROJANZOO: Everything you ever wanted to know about neural backdoors (but were afraid to ask)
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.
Dataset List
Model List
Attack List
Defense List
Note: This repository is also maintained to cover the implementation of
our kdd 2020 paper AdvMind: Inferring Adversary Intent of Black-Box Attacks
and ccs 2020 paper A Tale of Evil Twins: Adversarial Inputs versus Poisoned Models.
Screenshot
Features
- Colorful and verbose output!
Note: enable with
--color
for color and--verbose
for verbose.
To open an interactive window with color, usepython - --color
- Modular design (plug and play)
- Good code linting support (this package requires
python>=3.9
) - Register your own module to the library.
- Native Pytorch Output
trojanzoo
andtrojanvision
provides API to generate raw pytorch instances, which makes it flexible to work with nativepytorch
and other 3rd party libraries.trojanzoo.datasets.DataSet
can generatetorch.utils.data.Dataset
andtorch.utils.data.DataLoader
trojanzoo.models.Model
attribute_model
istorch.nn.Module
, attributemodel
istorch.nn.DataParallel
Specifically,trojanvision.datasets.ImageSet
can generatetorchvision.datasets.VisionDataset
,trojanvision.datasets.ImageFolder
can generatetorchvision.datasets.ImageFolder
- Enable pytorch native AMP(Automatic Mixed Precision) with
--amp
for training - Flexible Configuration Files
- Good help information to check arguments. (
-h
or--help
) - Detailed and well-organized
summary()
for each module.
Installation
pip install trojanzoo
- (todo)
conda install trojanzoo
docker pull local0state/trojanzoo
- (RECOMMEND)
python setup.py develop
This could install the github repo package but avoid copying files to
site_packages
, so that during developingTrojanZoo
, you don't need moveexample
files to the root directory of repo.
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
-
Train a model:
e.g.ResNet18
onCIFAR10
with 95% Accpython ./examples/train.py --color --verbose 1 --dataset cifar10 --model resnet18_comp --lr_scheduler --cutout --grad_clip 5.0 --save
-
Test backdoor attack (e.g., BadNet):
e.g.BadNet
withResNet18
onCIFAR10
python ./examples/backdoor_attack.py --color --verbose 1 --pretrain --validate_interval 1 --dataset cifar10 --model resnet18_comp --attack badnet --random_init --epoch 50 --lr 0.01 --save
-
Test backdoor defense (e.g., Neural Cleanse):
e.g.Neural Cleanse
againstBadNet
python ./examples/backdoor_defense.py --color --verbose 1 --pretrain --validate_interval 1 --dataset cifar10 --model resnet18_comp --attack badnet --defense neural_cleanse --random_init --epoch 50 --lr 0.01
IMC
python ./examples/backdoor_attack.py --color --verbose 1 --pretrain --validate_interval 1 --dataset cifar10 --model resnet18_comp --attack imc --random_init --epoch 50 --lr 0.01 --save
AdvMind
(with attack adaptive
and model adaptive
)
python ./examples/adv_defense.py --color --verbose 1 --pretrain --validate_interval 1 --dataset cifar10 --model resnet18_comp --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 overwritetrojanzoo
- 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. - User Default: Not decided yet.
(Enable it in the code
trojanzoo/configs/__init__.py
,trojanvision/configs/__init__.py
) - Workspace Default:
/configs/trojanzoo/
,/configs/trojanvision/
- Custom Config:
--config [config location]
- 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:
- CMD modules:
--verbose 1
- Colorful output:
--color
- tqdm progress bar:
--tqdm
- Check command-line argument usage:
--help
- AdvMind verbose information:
--output [number]
Use your DIY Dataset/Model/Attack/Defense
- 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
) - Register your DIY class in
trojanvision
Example:
trojanvision.attacks.class_dict[attack_name]=AttackClass
- Create your config files if necessary.
No need to modify any codes. Just directly add{attack_name}.yml
(.json
) in the config directory. - Good to go!
Todo List
- Sphinx Docs
- Unit test
License
TrojanZoo has a GPL-style license, as found in the LICENSE file.
Cite our paper
@InProceedings{pang2020trojanzoo,
title={TROJANZOO: Everything you ever wanted to know about neural backdoors (but were afraid to ask)},
author={Ren Pang and Zheng Zhang and Xiangshan Gao and Zhaohan Xi and Shouling Ji and Peng Cheng and Ting Wang},
year={2020},
booktitle={arXiv Preprint},
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for trojanzoo-1.0.10-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | ed6ddd7ce9f67bbb9a6e60aa3ed53a8c6fd2717bc12ccd5608f524027c2c617f |
|
MD5 | 8ffe6b240b40f7224c5092a5e22f4b6a |
|
BLAKE2b-256 | 21fc1de19026586bc4cca5a5c035df9c2e42dc5d989d25470278c0b1087a82af |