Please refer to: BasicReadMe
There are two version of our implementation. One is based on VP-SDE proposed by Song et al. (2021). This version is not recommended, and has been depreciated. For this, you need to download:
-
CIFAR10 unconditional diffusion model for DiffPure. Put it into ./resources/checkpoints/DiffPure/32x32_diffusion.pth
-
CIFAR10 WideResNet-70-16-dropout(discriminative classifier used in DiffPure). Put it into ./resources/checkpoints/models/WideResNet_70_16_dropout.pt
-
Conditional diffusion model for diffusion classifier: We will share our checkpoints soon. Now you can train it by yourself.
-
ImageNet unconditional diffusion model for DiffPure. Put it into ./resources/checkpoints/DiffPure/256x256_diffusion_uncond.pt
Another is based on Karras et al. (2022), a.k.a. EDM. We strongly recommend you to use this implementation. For this, you need to download:
- EDM checkpoints. Put it into ./resources/checkpoints/models/EDM/
- CIFAR10 WideResNet-70-16-dropout(discriminative classifier used in DiffPure). Put it into ./resources/checkpoints/models/WideResNet_70_16_dropout.pt
We provide some demos in the root. We believe it's very clear and you will know how to use DC/RDC/LM after playing with the example code. Hence, we didn't add more explanation in this README.
LM.py: Likelihood Maximization with WRN-70-2.
DCTK.py: Diffusion Classifier on CIFAR10.
RDCN+TK.py: Robust Diffusion Classifier on CIFAR10.
All experiments codes are in './experiments/'.
DiffAttack.py: Attack DiffPure.
DiffusionClassifierN+T: Multihead diffusion with likelihood maximization.
DiffusionClassifierTK: Off-the-shelf diffusion.
DiffusionClassifierN+TK: Off-the-shelf diffusion with likelihood maximization.
DiffusionAsClassifier: Depreciated. Test robustness of diffusion classifier (VP-SDE) under AutoAttack+BPDA/Lagrange/DirectDifferentiate
DiffusionMaximizer: Likelihood maximizer. A new diffusion purification method we proposed. See Sec 3.4 in our paper for detail. Could be combined with discriminative classifier.
DirectAttack: Direct differentiate through likelihood maximization.
ObfuscatedGradient: Measure the cosine similarity between the gradient of diffusion classifier and DiffPure. See Sec 4.4 in our paper for detail.
OptimalDiffusionClassifier: See Sec 3.3 for detail.
stadv: Measure the robustness under STAdv attack.
cifar100: Experiment on cifar100
Please cite us:
@article{chen2023robust,
title={Robust Classification via a Single Diffusion Model},
author={Chen, Huanran and Dong, Yinpeng and Wang, Zhengyi and Yang, Xiao and Duan, Chengqi and Su, Hang and Zhu, Jun},
journal={arXiv preprint arXiv:2305.15241},
year={2023}
}
If you have any question, you can contact us by:
Github issue.
Email: [email protected] (Recommended), [email protected] (Not recommended, this email will be banned 1 years later)