├── .gitignore ├── DCTK-NoShareNoise.py ├── DCTK.py ├── Intro.md ├── LM.py ├── RDCN+TK.py ├── README.md ├── attacks ├── AdversarialInput │ ├── AdversarialInputBase.py │ ├── AttentionConsistency.py │ ├── BIM.py │ ├── CommonWeakness.py │ ├── DI.py │ ├── DiffusionAttack.py │ ├── FGSM.py │ ├── MI.py │ ├── NAttack.py │ ├── Optim.py │ ├── PGD.py │ ├── PI.py │ ├── SVRE.py │ ├── SharpnessAware.py │ ├── SpectrumSimulationAttack.py │ ├── TI.py │ ├── VMI.py │ ├── __init__.py │ └── utils.py ├── __init__.py ├── autoattack │ ├── __init__.py │ ├── autoattack.py │ └── autopgd_base.py ├── perturbation │ ├── CommonWeakness.py │ ├── ParallelAttacker.py │ ├── PerturbationObject.py │ ├── SequentialAttacker.py │ ├── __init__.py │ └── utils.py ├── stadv_eot │ ├── __init__.py │ ├── attacks.py │ └── recoloradv │ │ ├── LICENSE_RECOLORADV │ │ ├── __init__.py │ │ ├── color_spaces.py │ │ ├── color_transformers.py │ │ ├── mister_ed │ │ ├── README.md │ │ ├── __init__.py │ │ ├── adversarial_attacks.py │ │ ├── adversarial_perturbations.py │ │ ├── adversarial_training.py │ │ ├── config.py │ │ ├── loss_functions.py │ │ ├── spatial_transformers.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ ├── checkpoints.py │ │ │ ├── discretization.py │ │ │ ├── image_utils.py │ │ │ ├── pytorch_ssim.py │ │ │ └── pytorch_utils.py │ │ ├── norms.py │ │ ├── perturbations.py │ │ └── utils.py └── utils.py ├── data ├── CIFAR10C.py ├── ImageNet.py ├── NIPS17.py ├── PACS.py ├── __init__.py ├── cifar.py ├── mnist.py ├── someset.py └── utils.py ├── defenses ├── HighLevelGuidedDenoiser │ ├── DYPUnet.py │ ├── Denoiser.py │ ├── __init__.py │ ├── dataset.py │ ├── defense.py │ ├── get_HGD.py │ ├── inception.py │ ├── inceptionresnet.py │ ├── inceptionresnetv2.py │ ├── inres.py │ ├── res152_wide.py │ ├── resnet.py │ ├── resnext.py │ ├── resnext101.py │ ├── resnext_features │ │ ├── __init__.py │ │ ├── resnext101_32x4d_features.py │ │ └── resnext101_64x4d_features.py │ └── v3.py ├── NeuralRepresentationPurifier │ ├── NRP.py │ ├── __init__.py │ ├── modules │ │ └── module_util.py │ └── networks.py ├── PurificationDefenses │ ├── DiffPure │ │ ├── DiffPure │ │ │ ├── DiffusionLikelihoodMaximizer.py │ │ │ ├── RS.py │ │ │ ├── __init__.py │ │ │ └── diffpure.py │ │ ├── DiffusionClassifier │ │ │ ├── DiffusionClassifier.py │ │ │ ├── DiffusionClassifierBase.py │ │ │ ├── DiffusionClassifierImageNet.py │ │ │ ├── DiffusionUseLessClassifier.py │ │ │ ├── EDMDC.py │ │ │ ├── OptimalDiffusionClassifier.py │ │ │ ├── SBGC.py │ │ │ ├── __init__.py │ │ │ └── utils.py │ │ ├── LikelihoodMaximization │ │ │ ├── EDMLM.py │ │ │ ├── LikelihoodMaximizationBase.py │ │ │ ├── __init__.py │ │ │ └── utils.py │ │ ├── __init__.py │ │ ├── losses │ │ │ ├── EDMLosses.py │ │ │ ├── __init__.py │ │ │ └── basic.py │ │ ├── model.py │ │ ├── sampler │ │ │ ├── BaseSampler.py │ │ │ ├── PreconditionChanger.py │ │ │ ├── __init__.py │ │ │ ├── ddim.py │ │ │ ├── edm.py │ │ │ └── sde.py │ │ └── utils.py │ ├── PurificationDefense.py │ └── __init__.py ├── RandomizedSmoothing │ ├── RSModels.py │ ├── __init__.py │ └── core.py ├── Transformations │ ├── BitDepthReduction.py │ ├── JPEGCompression.py │ ├── Randomization.py │ ├── TTE.py │ └── __init__.py └── __init__.py ├── experiments ├── CIFAR10C.py ├── DiffAttack.py ├── DiffusionAsClassifier.py ├── DiffusionAsClassifierN+T.py ├── DiffusionAsClassifierN+TK.py ├── DiffusionAsClassifierT.py ├── DiffusionAsClassifierTK.py ├── DiffusionMaximizer.py ├── DirectAttack.py ├── LikelihoodMaximization.py ├── ObfuscateGradient │ ├── ObfuscateGradient.py │ ├── ObfuscateGradient_likelihoodmaximizer.py │ └── Obfuscate_N+TK.py ├── OptimaDiffusionClassifier.py ├── RestrictedImageNet │ ├── ATImageNet.py │ ├── DiffAttackImageNet.py │ ├── EDMDCImageNet.py │ ├── EDMLMImageNet.py │ ├── EDMRDCImageNet.py │ ├── GuidedDiffusionClassifier.py │ └── TTEImageNet.py ├── cifar100 │ ├── CIFAR100DCNaN.py │ └── CIFAR100DiffPure.py └── stadv.py ├── models ├── BaseNormModel.py ├── RobustBench │ ├── __init__.py │ ├── cifar10.py │ └── imagenet.py ├── SmallResolutionModel │ ├── ShuffleNetv1.py │ ├── ShuffleNetv2.py │ ├── __init__.py │ ├── alexnet.py │ ├── cifar10_resnet.py │ ├── mobilenetv2.py │ ├── resnet.py │ ├── resnet_imagenet.py │ ├── resnetv2.py │ ├── rs_cifar_models.py │ ├── utils.py │ ├── vgg.py │ └── wrn.py ├── __init__.py ├── tf_models │ ├── DYPmodels.py │ ├── __init__.py │ ├── tf_adv_inception_v3.py │ ├── tf_ens3_adv_inc_v3.py │ ├── tf_ens4_adv_inc_v3.py │ ├── tf_ens_adv_inc_res_v2.py │ ├── tf_inc_res_v2.py │ ├── tf_inception_v3.py │ ├── tf_inception_v4.py │ ├── tf_resnet_v2_101.py │ ├── tf_resnet_v2_152.py │ └── tf_resnet_v2_50.py └── unets │ ├── DiffuserUNets.py │ ├── EDM │ ├── __init__.py │ ├── dnnlib │ │ ├── __init__.py │ │ └── util.py │ ├── edm_nets.py │ ├── get_edm_nets.py │ └── torch_utils │ │ ├── __init__.py │ │ ├── distributed.py │ │ ├── misc.py │ │ ├── persistence.py │ │ └── training_stats.py │ ├── __init__.py │ ├── guided_diffusion │ ├── LICENSE_GUIDED_DIFFUSION │ ├── __init__.py │ ├── classifier.py │ ├── config.py │ ├── default_unet.py │ ├── dist_util.py │ ├── fp16_util.py │ ├── gaussian_diffusion.py │ ├── image_datasets.py │ ├── logger.py │ ├── losses.py │ ├── nn.py │ ├── resample.py │ ├── respace.py │ ├── script_util.py │ ├── train_util.py │ └── unet.py │ ├── improved_diffusion │ ├── __init__.py │ ├── dist_util.py │ ├── fp16_util.py │ ├── gaussian_diffusion.py │ ├── image_datasets.py │ ├── logger.py │ ├── losses.py │ ├── networks │ │ ├── __init__.py │ │ ├── lenet.py │ │ ├── resnet.py │ │ ├── vggnet.py │ │ └── wide_resnet.py │ ├── nn.py │ ├── resample.py │ ├── respace.py │ ├── script_util.py │ ├── train_util.py │ └── unet.py │ └── score_sde │ ├── LICENSE_SCORE_SDE │ ├── MultiHeadRDC │ ├── DiffusionClassifier.py │ ├── __init__.py │ └── utils.py │ ├── __init__.py │ ├── cifar10.yml │ ├── losses.py │ ├── models │ ├── __init__.py │ ├── ddpm.py │ ├── ema.py │ ├── layers.py │ ├── layerspp.py │ ├── ncsnpp.py │ ├── ncsnpp_multihead.py │ ├── ncsnv2.py │ ├── normalization.py │ ├── up_or_down_sampling.py │ └── utils.py │ ├── op │ ├── __init__.py │ ├── fused_act.py │ ├── fused_bias_act.cpp │ ├── fused_bias_act_kernel.cu │ ├── upfirdn2d.cpp │ ├── upfirdn2d.py │ └── upfirdn2d_kernel.cu │ ├── sampling.py │ ├── sde_lib.py │ └── unets.py ├── optimizer ├── ALRS.py ├── PGD.py ├── __init__.py ├── default.py └── losses │ ├── CKA.py │ └── __init__.py ├── requirements.txt ├── tester ├── AutoAttack.py ├── CertifyRobustness.py ├── TestAcc.py ├── TransferAttackAcc.py ├── __init__.py ├── generation.py └── utils.py └── utils ├── HRNet ├── Changer.py ├── HRBatchNorm.py ├── HRConv.py ├── HRLinear.py ├── HRReLU.py └── __init__.py ├── ImageHandling.py ├── __init__.py ├── gpu.py ├── plot ├── ColorUtils.py ├── CommonFigures.py ├── Landscape4Input.py ├── Landscape4Model.py └── __init__.py ├── saver.py ├── seed.py └── types.py /.gitignore: -------------------------------------------------------------------------------- 1 | ### Python template 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .resources 12 | .resources/* 13 | .idea 14 | .datasets 15 | .Python 16 | build/ 17 | develop-eggs/ 18 | dist/ 19 | downloads/ 20 | eggs/ 21 | .eggs/ 22 | lib/ 23 | lib64/ 24 | parts/ 25 | sdist/ 26 | var/ 27 | wheels/ 28 | share/python-wheels/ 29 | *.egg-info/ 30 | .installed.cfg 31 | *.egg 32 | MANIFEST 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .nox/ 48 | .coverage 49 | .coverage.* 50 | .cache 51 | nosetests.xml 52 | coverage.xml 53 | *.cover 54 | *.py,cover 55 | .hypothesis/ 56 | .pytest_cache/ 57 | cover/ 58 | 59 | # Translations 60 | *.mo 61 | *.pot 62 | 63 | # Django stuff: 64 | *.log 65 | local_settings.py 66 | db.sqlite3 67 | db.sqlite3-journal 68 | 69 | # Flask stuff: 70 | instance/ 71 | .webassets-cache 72 | 73 | # Scrapy stuff: 74 | .scrapy 75 | 76 | # Sphinx documentation 77 | docs/_build/ 78 | 79 | # PyBuilder 80 | .pybuilder/ 81 | target/ 82 | 83 | # Jupyter Notebook 84 | .ipynb_checkpoints 85 | 86 | # IPython 87 | profile_default/ 88 | ipython_config.py 89 | 90 | # pyenv 91 | # For a library or package, you might want to ignore these files since the code is 92 | # intended to run in multiple environments; otherwise, check them in: 93 | # .python-version 94 | 95 | # pipenv 96 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 97 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 98 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 99 | # install all needed dependencies. 100 | #Pipfile.lock 101 | 102 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 103 | __pypackages__/ 104 | 105 | # Celery stuff 106 | celerybeat-schedule 107 | celerybeat.pid 108 | 109 | # SageMath parsed files 110 | *.sage.py 111 | 112 | # Environments 113 | .env 114 | .venv 115 | env/ 116 | venv/ 117 | ENV/ 118 | env.bak/ 119 | venv.bak/ 120 | 121 | # Spyder project settings 122 | .spyderproject 123 | .spyproject 124 | 125 | # Rope project settings 126 | .ropeproject 127 | 128 | # mkdocs documentation 129 | /site 130 | 131 | # mypy 132 | .mypy_cache/ 133 | .dmypy.json 134 | dmypy.json 135 | 136 | # Pyre type checker 137 | .pyre/ 138 | 139 | # pytype static type analyzer 140 | .pytype/ 141 | 142 | # Cython debug symbols 143 | cython_debug/ 144 | 145 | /resources 146 | 147 | -------------------------------------------------------------------------------- /DCTK-NoShareNoise.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 7 | 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument("--begin", type=int, default=0) 11 | parser.add_argument("--end", type=int, default=128) 12 | args = parser.parse_args() 13 | begin, end = args.begin, args.end 14 | 15 | model = get_edm_cifar_cond().cuda() 16 | test_loader = get_CIFAR10_test(batch_size=1) 17 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 18 | 19 | dc = EDMEulerIntegralWraped(unet=model, timesteps=torch.linspace(1e-4, 3, 1001)) 20 | test_acc(dc, test_loader) 21 | test_apgd_dlr_acc(dc, loader=test_loader) 22 | -------------------------------------------------------------------------------- /DCTK.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, DiffusionClassifierSingleHeadBaseWraped 7 | 8 | """ 9 | Example code of Diffusion Classifier! 10 | """ 11 | 12 | parser = argparse.ArgumentParser() 13 | parser.add_argument("--begin", type=int, default=0) 14 | parser.add_argument("--end", type=int, default=128) 15 | parser.add_argument("--steps", type=int, default=126) 16 | args = parser.parse_args() 17 | begin, end, steps = args.begin, args.end, args.steps 18 | 19 | model = get_edm_cifar_cond().cuda() 20 | test_loader = get_CIFAR10_test(batch_size=1) 21 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 22 | 23 | dc = EDMEulerIntegralDC(unet=model, timesteps=torch.linspace(1e-4, 3, steps)) 24 | dc.share_noise = True 25 | dc_wraped = DiffusionClassifierSingleHeadBaseWraped(dc) 26 | test_acc(dc_wraped, test_loader) 27 | test_apgd_dlr_acc(dc_wraped, loader=test_loader) 28 | -------------------------------------------------------------------------------- /LM.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from models import WideResNet_70_16_dropout 4 | from data import get_CIFAR10_test 5 | from tester import test_acc, test_apgd_dlr_acc 6 | import argparse 7 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralLM, \ 8 | diffusion_likelihood_maximizer_defense 9 | 10 | """ 11 | Example code of Likelihood Maximization! 12 | """ 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument("--begin", type=int, default=0) 16 | parser.add_argument("--end", type=int, default=512) 17 | args = parser.parse_args() 18 | begin, end = args.begin, args.end 19 | 20 | model = get_edm_cifar_cond().cuda() 21 | test_loader = get_CIFAR10_test(batch_size=1) 22 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 23 | 24 | uncond_edm_net = get_edm_cifar_uncond() 25 | uncond_edm_dc = EDMEulerIntegralDC(unet=uncond_edm_net) 26 | lm = EDMEulerIntegralLM(uncond_edm_dc) 27 | 28 | lm_wrn = diffusion_likelihood_maximizer_defense( 29 | WideResNet_70_16_dropout(), 30 | lm.likelihood_maximization_T1, 31 | ) 32 | 33 | test_acc(lm_wrn, test_loader) 34 | test_apgd_dlr_acc(lm_wrn, loader=test_loader) 35 | -------------------------------------------------------------------------------- /RDCN+TK.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralLM, \ 7 | diffusion_likelihood_maximizer_defense 8 | 9 | 10 | """ 11 | Example code of Robust Diffusion Classifier! 12 | """ 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument("--begin", type=int, default=0) 16 | parser.add_argument("--end", type=int, default=512) 17 | parser.add_argument("--steps", type=int, default=126) 18 | args = parser.parse_args() 19 | begin, end, steps = args.begin, args.end, args.steps 20 | 21 | model = get_edm_cifar_cond().cuda() 22 | test_loader = get_CIFAR10_test(batch_size=1) 23 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 24 | 25 | dc = EDMEulerIntegralDC(unet=model, timesteps=torch.linspace(1e-4, 3, steps)) 26 | dc.share_noise = True 27 | 28 | uncond_edm_net = get_edm_cifar_uncond() 29 | uncond_edm_dc = EDMEulerIntegralDC(unet=uncond_edm_net) 30 | lm = EDMEulerIntegralLM(uncond_edm_dc) 31 | 32 | rdc = diffusion_likelihood_maximizer_defense( 33 | dc, 34 | lm.likelihood_maximization_T1, 35 | ) 36 | 37 | test_acc(rdc, test_loader) 38 | test_apgd_dlr_acc(rdc, loader=test_loader) 39 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/BIM.py: -------------------------------------------------------------------------------- 1 | from .PGD import PGD 2 | 3 | 4 | class BIM(PGD): 5 | def __init__(self, *args, **kwargs): 6 | kwargs['random_start'] = False 7 | super(BIM, self).__init__(*args, **kwargs) 8 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/DI.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks.utils import * 3 | from torch import nn 4 | from typing import Callable, List 5 | from .AdversarialInputBase import AdversarialInputAttacker 6 | from torchvision import transforms 7 | 8 | 9 | class DI_MI_FGSM(AdversarialInputAttacker): 10 | ''' 11 | DI-FGSM is not using data augmentation to increase data for optimizing perturbations. 12 | DI-FGSM actually is using differentiable data augmentations, 13 | and this data augmentation can be viewed as a part of model(from SI-FGSM) 14 | ''' 15 | 16 | def __init__(self, 17 | model: List[nn.Module], 18 | total_step: int = 10, 19 | random_start: bool = False, 20 | step_size: float = 16 / 255 / 10, 21 | criterion: Callable = nn.CrossEntropyLoss(), 22 | targeted_attack=False, 23 | mu: float = 1, 24 | *args, **kwargs 25 | ): 26 | super(DI_MI_FGSM, self).__init__(model, *args, **kwargs) 27 | self.random_start = random_start 28 | self.total_step = total_step 29 | self.step_size = step_size 30 | self.criterion = criterion 31 | self.targerted_attack = targeted_attack 32 | self.mu = mu 33 | self.aug_policy = transforms.Compose([ 34 | transforms.RandomCrop((224, 224), padding=224 - int(224 * 0.9)), 35 | ]) 36 | self.init() 37 | 38 | def perturb(self, x): 39 | x = x + (torch.rand_like(x) - 0.5) * 2 * self.epsilon 40 | x = clamp(x) 41 | return x 42 | 43 | def attack(self, x, y, ): 44 | N = x.shape[0] 45 | original_x = x.clone() 46 | momentum = torch.zeros_like(x) 47 | if self.random_start: 48 | x = self.perturb(x) 49 | 50 | for _ in range(self.total_step): 51 | x.requires_grad = True 52 | aug_x = self.aug_policy(x) 53 | logit = 0 54 | for model in self.models: 55 | logit += model(aug_x.to(model.device)).to(x.device) 56 | loss = self.criterion(logit, y) 57 | loss.backward() 58 | grad = x.grad 59 | x.requires_grad = False 60 | # update 61 | if self.targerted_attack: 62 | momentum = self.mu * momentum - grad / torch.norm(grad.reshape(N, -1), p=1, dim=1).view(N, 1, 1, 1) 63 | x += self.step_size * momentum.sign() 64 | else: 65 | momentum = self.mu * momentum + grad / torch.norm(grad.reshape(N, -1), p=1, dim=1).view(N, 1, 1, 1) 66 | x += self.step_size * momentum.sign() 67 | x = clamp(x) 68 | x = clamp(x, original_x - self.epsilon, original_x + self.epsilon) 69 | 70 | return x 71 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/FGSM.py: -------------------------------------------------------------------------------- 1 | from .PGD import PGD 2 | 3 | 4 | class FGSM(PGD): 5 | def __init__(self, *args, epsilon=16 / 255, random_start=False, **kwargs): 6 | kwargs['total_step'] = 1 7 | kwargs['random_start'] = random_start 8 | kwargs['epsilon'] = epsilon 9 | kwargs['step_size'] = epsilon 10 | super(FGSM, self).__init__(*args, **kwargs) 11 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/MI.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks.utils import * 3 | from torch import nn 4 | from typing import Callable, List 5 | from .AdversarialInputBase import AdversarialInputAttacker 6 | 7 | 8 | class MI_FGSM(AdversarialInputAttacker): 9 | def __init__(self, 10 | model: List[nn.Module], 11 | total_step: int = 10, random_start: bool = False, 12 | step_size: float = 16 / 255 / 10, 13 | criterion: Callable = nn.CrossEntropyLoss(), 14 | targeted_attack=False, 15 | mu: float = 1, 16 | *args, **kwargs 17 | ): 18 | self.random_start = random_start 19 | self.total_step = total_step 20 | self.step_size = step_size 21 | self.criterion = criterion 22 | self.targerted_attack = targeted_attack 23 | self.mu = mu 24 | super(MI_FGSM, self).__init__(model, *args, **kwargs) 25 | 26 | def perturb(self, x): 27 | x = x + (torch.rand_like(x) - 0.5) * 2 * self.epsilon 28 | x = clamp(x) 29 | return x 30 | 31 | def attack(self, x, y, ): 32 | N = x.shape[0] 33 | original_x = x.clone() 34 | momentum = torch.zeros_like(x) 35 | if self.random_start: 36 | x = self.perturb(x) 37 | 38 | for _ in range(self.total_step): 39 | x.requires_grad = True 40 | logit = 0 41 | for model in self.models: 42 | logit += model(x.to(model.device)).to(x.device) 43 | loss = self.criterion(logit, y) 44 | loss.backward() 45 | grad = x.grad 46 | x.requires_grad = False 47 | # update 48 | if self.targerted_attack: 49 | momentum = self.mu * momentum - grad / torch.norm(grad.reshape(N, -1), p=1, dim=1).view(N, 1, 1, 1) 50 | x += self.step_size * momentum.sign() 51 | else: 52 | momentum = self.mu * momentum + grad / torch.norm(grad.reshape(N, -1), p=1, dim=1).view(N, 1, 1, 1) 53 | x += self.step_size * momentum.sign() 54 | x = self.clamp(x, original_x) 55 | return x 56 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/NAttack.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks.utils import * 3 | from torch import nn 4 | from typing import Callable, List 5 | from .AdversarialInputBase import AdversarialInputAttacker 6 | 7 | __all__ = ['NAttack'] 8 | 9 | 10 | class NAttack(AdversarialInputAttacker): 11 | def __init__(self, model: List[nn.Module], 12 | total_step: int = 600, 13 | step_size: float = 0.008, 14 | batch_size=300, 15 | sigma=0.1, 16 | ): 17 | self.total_step = total_step 18 | self.step_size = step_size 19 | self.batch_size = batch_size 20 | super(NAttack, self).__init__(model, *args, **kwargs) 21 | self.sigma = sigma 22 | 23 | def perturb(self, x): 24 | x = x + (torch.rand_like(x) - 0.5) * 2 * self.epsilon 25 | x = clamp(x) 26 | return x 27 | 28 | def g(self, x): 29 | return (torch.tanh(x) + 1) / 2 30 | 31 | def proj(self, mu, noise): 32 | return clamp(clamp(noise, mu - self.epsilon, mu + self.epsilon)) 33 | 34 | @torch.no_grad() 35 | def attack(self, x, y, ): 36 | assert x.shape[0] == 1, 'now only support batch size = 1' 37 | x.requires_grad = False 38 | N, C, H, D = x.shape 39 | original_x = x.clone() 40 | mu = torch.arctan(2 * x - 1) 41 | mu_min = torch.arctan(2 * (x - self.epsilon) - 1) 42 | mu_max = torch.arctan(2 * (x + self.epsilon) - 1) 43 | for _ in range(self.total_step): 44 | epsilons = torch.randn(self.batch_size, C, H, D, device=self.device) 45 | samples = mu + self.sigma * epsilons 46 | samples = self.g(samples) 47 | samples = self.proj(original_x, samples) 48 | z_scores = [] 49 | for model in self.models: 50 | out = model(samples) 51 | pre = torch.max(out, dim=1)[1] 52 | mask = (pre != y).float() 53 | if torch.sum(mask) > 0: 54 | index = torch.max(mask, dim=0)[1] 55 | result = samples[index].unsqueeze(0) 56 | return result 57 | f = out[:, y] # batch_size 58 | z_score = (f - torch.mean(f)) / (torch.std(f) + 1e-7) 59 | z_scores.append(z_score) 60 | z_scores = torch.stack(z_scores).mean(0) # batch_size 61 | mu = mu - self.step_size / (self.batch_size * self.sigma) * (z_scores.view(-1, 1, 1, 1) * epsilons).sum(0) 62 | mu = clamp(mu, min_value=mu_min, max_value=mu_max) 63 | result = self.proj(original_x, self.g(mu)) 64 | return result 65 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/Optim.py: -------------------------------------------------------------------------------- 1 | ''' 2 | PGD: Projected Gradient Descent 3 | ''' 4 | 5 | import torch 6 | from torch import nn 7 | from typing import Callable, List 8 | from attacks.utils import * 9 | from .AdversarialInputBase import AdversarialInputAttacker 10 | from torch.optim import Optimizer, Adam 11 | 12 | 13 | class OptimizerAttacker(AdversarialInputAttacker): 14 | def __init__(self, 15 | model: List[nn.Module], 16 | optimizer: Callable = lambda x: Adam([x], lr=0.1), 17 | total_step: int = 10, 18 | random_start: bool = False, 19 | criterion: Callable = nn.CrossEntropyLoss(), 20 | *args, 21 | **kwargs 22 | ): 23 | self.random_start = random_start 24 | self.criterion = criterion 25 | self.optimizer = optimizer 26 | self.total_step = total_step 27 | super(OptimizerAttacker, self).__init__(model, *args, **kwargs) 28 | 29 | def perturb(self, x): 30 | x = x + (torch.rand_like(x) - 0.5) * 2 * self.epsilon 31 | x = clamp(x) 32 | return x 33 | 34 | def attack(self, x, y, ): 35 | x.requires_grad_() 36 | optimizer = self.optimizer(x) 37 | original_x = x.clone() 38 | if self.random_start: 39 | x = self.perturb(x) 40 | 41 | for _ in range(self.total_step): 42 | logit = 0 43 | for model in self.models: 44 | logit += model(x.to(model.device)).to(x.device) 45 | loss = self.criterion(logit, y) 46 | optimizer.zero_grad() 47 | loss.backward() 48 | optimizer.step() 49 | self.clamp(x, original_x, inplace=True) 50 | x.requires_grad_(False) 51 | return x 52 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/PGD.py: -------------------------------------------------------------------------------- 1 | """ 2 | PGD: Projected Gradient Descent 3 | """ 4 | 5 | import torch 6 | from torch import nn, Tensor 7 | from typing import Callable, List 8 | from attacks.utils import * 9 | from .AdversarialInputBase import AdversarialInputAttacker 10 | 11 | 12 | class PGD(AdversarialInputAttacker): 13 | def __init__( 14 | self, 15 | model: List[nn.Module], 16 | total_step: int = 10, 17 | random_start: bool = True, 18 | step_size: float = 16 / 255 / 10, 19 | eot_step: int = 1, 20 | eot_batch_size: int = 1024, 21 | criterion: Callable = nn.CrossEntropyLoss(), 22 | targeted_attack=False, 23 | *args, 24 | **kwargs, 25 | ): 26 | self.random_start = random_start 27 | self.total_step = total_step 28 | self.step_size = step_size 29 | self.eot_step = eot_step 30 | self.eot_batch_size = eot_batch_size 31 | self.criterion = criterion 32 | self.targerted_attack = targeted_attack 33 | super(PGD, self).__init__(model, *args, **kwargs) 34 | 35 | def perturb(self, x): 36 | x = x + (torch.rand_like(x) - 0.5) * 2 * self.epsilon 37 | x = clamp(x) 38 | return x 39 | 40 | def attack(self, x: Tensor, y: Tensor): 41 | assert len(x.shape) == 4, "input should have size B, C, H, D" 42 | B, C, H, D = x.shape 43 | original_x = x.clone() 44 | if self.random_start: 45 | x = self.perturb(x) 46 | 47 | for _ in range(self.total_step): 48 | x.requires_grad = True 49 | eot_xs = x.repeat(self.eot_step, 1, 1, 1).split(self.eot_batch_size * B, dim=0) 50 | for eot_x in eot_xs: 51 | logit = 0 52 | for model in self.models: 53 | logit += model(eot_x.to(model.device)).to(x.device) 54 | loss = self.criterion(logit, y.repeat(eot_x.shape[0] // y.shape[0])) 55 | loss.backward() 56 | grad = x.grad / self.eot_step 57 | x.requires_grad = False 58 | # update 59 | if self.targerted_attack: 60 | x = x - self.step_size * (grad.sign() if self.norm == "Linf" else self.normalize(grad)) 61 | else: 62 | x = x + self.step_size * (grad.sign() if self.norm == "Linf" else self.normalize(grad)) 63 | x = self.clamp(x, original_x) 64 | return x 65 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/TI.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks.utils import * 3 | from torch import nn 4 | from typing import Callable, List 5 | from .AdversarialInputBase import AdversarialInputAttacker 6 | import numpy as np 7 | from scipy import stats as st 8 | 9 | 10 | class MI_TI_FGSM(AdversarialInputAttacker): 11 | def __init__(self, model: List[nn.Module], 12 | total_step: int = 10, random_start: bool = False, 13 | step_size: float = 16 / 255 / 10, 14 | criterion: Callable = nn.CrossEntropyLoss(), 15 | targeted_attack=False, 16 | mu: float = 1, 17 | *args, **kwargs, 18 | ): 19 | self.random_start = random_start 20 | self.total_step = total_step 21 | self.step_size = step_size 22 | self.criterion = criterion 23 | self.targerted_attack = targeted_attack 24 | self.mu = mu 25 | super(MI_TI_FGSM, self).__init__(model, *args, **kwargs) 26 | self.conv = self.gkern().to(self.device) 27 | self.conv.requires_grad_(False) 28 | 29 | def perturb(self, x): 30 | x = x + (torch.rand_like(x) - 0.5) * 2 * self.epsilon 31 | x = clamp(x) 32 | return x 33 | 34 | def attack(self, x, y, ): 35 | N = x.shape[0] 36 | original_x = x.clone() 37 | momentum = torch.zeros_like(x) 38 | if self.random_start: 39 | x = self.perturb(x) 40 | 41 | for _ in range(self.total_step): 42 | x.requires_grad = True 43 | logit = 0 44 | for model in self.models: 45 | logit += model(x.to(model.device)).to(x.device) 46 | loss = self.criterion(logit, y) 47 | loss.backward() 48 | grad = x.grad 49 | x.requires_grad = False 50 | # update 51 | grad = self.conv(grad) 52 | if self.targerted_attack: 53 | momentum = self.mu * momentum - grad / torch.norm(grad.reshape(N, -1), p=1, dim=1).view(N, 1, 1, 1) 54 | x += self.step_size * momentum.sign() 55 | else: 56 | momentum = self.mu * momentum + grad / torch.norm(grad.reshape(N, -1), p=1, dim=1).view(N, 1, 1, 1) 57 | x += self.step_size * momentum.sign() 58 | x = clamp(x) 59 | x = clamp(x, original_x - self.epsilon, original_x + self.epsilon) 60 | 61 | return x 62 | 63 | @staticmethod 64 | def gkern(kernlen=15, nsig=3): 65 | """Returns a 2D Gaussian kernel array.""" 66 | x = np.linspace(-nsig, nsig, kernlen) 67 | kern1d = st.norm.pdf(x) 68 | kernel_raw = np.outer(kern1d, kern1d) 69 | kernel = kernel_raw / kernel_raw.sum() 70 | kernel = torch.tensor(kernel, dtype=torch.float) 71 | conv = nn.Conv2d(3, 3, kernel_size=kernlen, stride=1, padding=kernlen // 2, bias=False, groups=3) 72 | kernel = kernel.repeat(3, 1, 1).view(3, 1, kernlen, kernlen) 73 | conv.weight.data = kernel 74 | return conv 75 | 76 | 77 | if __name__ == '__main__': 78 | MI_TI_FGSM.gkern() 79 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/__init__.py: -------------------------------------------------------------------------------- 1 | from .CommonWeakness import * 2 | from .SharpnessAware import * 3 | from .MI import MI_FGSM 4 | from .FGSM import * 5 | from .BIM import * 6 | from .VMI import * 7 | from .TI import * 8 | from .Optim import * 9 | from .DI import DI_MI_FGSM 10 | from .SVRE import MI_SVRE 11 | from .NAttack import * 12 | from .PI import PI_FGSM 13 | from .SpectrumSimulationAttack import SpectrumSimulationAttack, SSA_CommonWeakness 14 | from .AttentionConsistency import MI_AttentionConsistency 15 | from .AdversarialInputBase import AdversarialInputAttacker, IdentityAttacker 16 | -------------------------------------------------------------------------------- /attacks/AdversarialInput/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def cosine_similarity(x: list): 5 | """ 6 | input a list of tensor with same shape. return the mean cosine_similarity 7 | """ 8 | x = torch.stack(x, dim=0) 9 | N = x.shape[0] 10 | x = x.reshape(N, -1) 11 | 12 | norm = torch.norm(x, p=2, dim=1) 13 | x /= norm.reshape(-1, 1) # N, D 14 | similarity = x @ x.T # N, N 15 | # matrix_heatmap(similarity.cpu().numpy()) 16 | mask = torch.triu(torch.ones(N, N, device=x.device), diagonal=0).to(torch.bool) 17 | similarity = similarity[mask] 18 | return torch.mean(similarity).item() 19 | 20 | 21 | def clamp(x: torch.tensor, min_value: float = 0, max_value: float = 1): 22 | return torch.clamp(x, min=min_value, max=max_value) 23 | -------------------------------------------------------------------------------- /attacks/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | all the implementation here, the range of images values is [0, 1] 3 | if a model need normalization, please adding the normalization part in the model, not in loader or attacker 4 | 5 | ''' 6 | 7 | from .perturbation import * 8 | from .AdversarialInput import * 9 | from .stadv_eot import * 10 | -------------------------------------------------------------------------------- /attacks/autoattack/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | AutoAttack are not compatible with gradient checkpoint! 3 | We fix this bug! 4 | """ 5 | from .autoattack import AutoAttack 6 | -------------------------------------------------------------------------------- /attacks/perturbation/CommonWeakness.py: -------------------------------------------------------------------------------- 1 | from .SequentialAttacker import SequentialAttacker 2 | from typing import Callable, List, Iterable 3 | from attacks.utils import * 4 | from .utils import cosine_similarity 5 | 6 | 7 | class CosineSimilarityEncourager(SequentialAttacker): 8 | def __init__(self, *args, **kwargs): 9 | super(CosineSimilarityEncourager, self).__init__(*args, **kwargs) 10 | if kwargs['outer_optimizer'] is not None: 11 | self.outer_optimizer = kwargs['outer_optimizer']([self.perturbation.perturbation]) 12 | else: 13 | self.outer_optimizer = None 14 | 15 | def attack(self, 16 | loader: DataLoader or Iterable, 17 | total_iter_step: int = 10, 18 | is_clamp=True): 19 | iter_step = 0 20 | while True: 21 | for x, y in loader: 22 | original_x = x.clone() 23 | self.begin_attack() 24 | for model in self.models: 25 | x = original_x + self.perturbation.perturbation 26 | if is_clamp: 27 | x = clamp(x) 28 | x = self.transform(x) 29 | loss = self.criterion(model(x.to(model.device)), y.to(model.device)) 30 | self.perturbation.zero_grad() 31 | loss.backward() 32 | self.grad_record.append(self.perturbation.grad()) 33 | self.perturbation.step() 34 | self.end_attack() 35 | iter_step += 1 36 | if iter_step > total_iter_step: 37 | self.perturbation.requires_grad(False) 38 | return self.perturbation 39 | 40 | @torch.no_grad() 41 | def begin_attack(self): 42 | self.original = self.perturbation.perturbation.clone() 43 | self.grad_record = [] 44 | 45 | @torch.no_grad() 46 | def end_attack(self, ksi=1): 47 | ''' 48 | theta: original_patch 49 | theta_hat: now patch in optimizer 50 | theta = theta + ksi*(theta_hat - theta), so: 51 | theta =(1-ksi )theta + ksi* theta_hat 52 | ''' 53 | patch = self.perturbation.perturbation 54 | if self.outer_optimizer is None: 55 | patch.mul_(ksi) 56 | patch.add_((1 - ksi) * self.original) 57 | else: 58 | fake_grad = - ksi * (patch - self.original) 59 | self.outer_optimizer.zero_grad() 60 | patch.mul_(0) 61 | patch.add_(self.original) 62 | patch.grad = fake_grad 63 | self.outer_optimizer.step() 64 | self.perturbation.clamp() 65 | 66 | grad_similarity = cosine_similarity(self.grad_record) 67 | del self.grad_record 68 | del self.original 69 | -------------------------------------------------------------------------------- /attacks/perturbation/ParallelAttacker.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, List, Iterable 2 | from torch import nn 3 | from attacks.utils import * 4 | from .PerturbationObject import Perturbation 5 | from .SequentialAttacker import SequentialAttacker 6 | 7 | 8 | class ParallelAttacker(SequentialAttacker): 9 | ''' 10 | please set your learning rate in optimizer 11 | set data augmentation in your loader. 12 | ''' 13 | 14 | def __init__(self, *args, **kwargs, 15 | ): 16 | super(ParallelAttacker, self).__init__(*args, **kwargs) 17 | 18 | def attack(self, 19 | loader: DataLoader or Iterable, 20 | total_iter_step: int = 10, 21 | is_clamp=True): 22 | iter_step = 0 23 | while True: 24 | for x, y in loader: 25 | x = x + self.perturbation.perturbation 26 | if is_clamp: 27 | x = clamp(x) 28 | x = self.transform(x) 29 | 30 | loss = 0 31 | for model in self.models: 32 | loss += self.criterion(model(x.to(model.device)), y.to(model.device)) 33 | self.perturbation.zero_grad() 34 | loss.backward() 35 | self.perturbation.step() 36 | 37 | iter_step += 1 38 | if iter_step > total_iter_step: 39 | self.perturbation.requires_grad(False) 40 | return self.perturbation 41 | -------------------------------------------------------------------------------- /attacks/perturbation/PerturbationObject.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.optim import Optimizer 3 | from typing import Callable, List, Iterable 4 | from attacks.utils import * 5 | from .utils import * 6 | 7 | 8 | class Perturbation(): 9 | def __init__(self, 10 | optimizer: Optimizer or Callable, 11 | perturbation_size: tuple = (3, 299, 299), 12 | device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), 13 | epsilon=16 / 255, ): 14 | self.perturbation = torch.zeros(perturbation_size, device=device) 15 | self.device = device 16 | self.optimizer = optimizer 17 | self.eps = epsilon 18 | 19 | def gaussian_init(self, is_clamp=True, scale=0.5, mean=0.5): 20 | self.perturbation = torch.randn_like(self.perturbation, device=self.device) * scale + mean 21 | if is_clamp: 22 | self.perturbation = clamp(self.perturbation) 23 | self.optimizer = self.optimizer([self.perturbation]) 24 | 25 | def uniform_init(self): 26 | self.perturbation = torch.rand_like(self.perturbation, device=self.device) 27 | self.optimizer = self.optimizer([self.perturbation]) 28 | 29 | def constant_init(self, constant=0): 30 | self.perturbation = torch.zeros_like(self.perturbation, device=self.device) + constant 31 | self.optimizer = self.optimizer([self.perturbation]) 32 | 33 | def requires_grad(self, requires_grad: bool = True): 34 | self.perturbation.requires_grad = requires_grad 35 | 36 | @torch.no_grad() 37 | def step(self): 38 | self.optimizer.step() 39 | self.perturbation.clamp_(min=-self.eps, max=self.eps) 40 | 41 | def clamp(self): 42 | self.perturbation.clamp_(min=-self.eps, max=self.eps) 43 | 44 | def zero_grad(self): 45 | self.optimizer.zero_grad() 46 | 47 | def grad(self): 48 | return self.perturbation.grad 49 | 50 | def assign_grad(self, grad): 51 | self.perturbation.grad = grad 52 | -------------------------------------------------------------------------------- /attacks/perturbation/SequentialAttacker.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, List, Iterable 2 | import torch 3 | from torch import nn 4 | from attacks.utils import * 5 | from .PerturbationObject import Perturbation 6 | from torch.nn import functional as F 7 | 8 | 9 | class SequentialAttacker(): 10 | ''' 11 | please set your learning rate in optimizer 12 | set data augmentation in your loader. 13 | ''' 14 | 15 | def __init__(self, 16 | models: List[nn.Module], 17 | perturbation: Perturbation, 18 | transformation: nn.Module = nn.Identity(), 19 | criterion: Callable = F.cross_entropy, 20 | **kwargs, 21 | ): 22 | self.perturbation = torch.randn 23 | self.models = models 24 | perturbation.requires_grad(True) 25 | self.perturbation = perturbation 26 | self.transform = transformation 27 | self.criterion = criterion 28 | self.init() 29 | 30 | def init(self): 31 | for i, model in enumerate(self.models): 32 | model.requires_grad_(False) 33 | model.to(torch.device(f'cuda:{i}')) 34 | model.device = torch.device(f'cuda:{i}') 35 | self.perturbation.requires_grad(True) 36 | 37 | def tensor_to_loader(self, x, y): 38 | return [(x, y)] 39 | 40 | def attack(self, 41 | loader: DataLoader or Iterable, 42 | total_iter_step: int = 10, 43 | is_clamp=True): 44 | iter_step = 0 45 | while True: 46 | for x, y in loader: 47 | original_x = x.clone() 48 | for model in self.models: 49 | x = original_x + self.perturbation.perturbation 50 | if is_clamp: 51 | x = clamp(x) 52 | x = self.transform(x) 53 | loss = self.criterion(model(x.to(model.device)), y.to(model.device)) 54 | self.perturbation.zero_grad() 55 | loss.backward() 56 | self.perturbation.step() 57 | 58 | iter_step += 1 59 | if iter_step > total_iter_step: 60 | self.perturbation.requires_grad(False) 61 | return self.perturbation 62 | 63 | def __call__(self, x, y, total_iter_step=20): 64 | with torch.no_grad(): 65 | self.perturbation.perturbation.mul_(0) 66 | self.perturbation.requires_grad(True) 67 | p = self.attack(self.tensor_to_loader(x, y), total_iter_step=total_iter_step) 68 | p.requires_grad(False) 69 | return x + p.perturbation 70 | -------------------------------------------------------------------------------- /attacks/perturbation/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Perturbation attackers are all ensemble attacker. 3 | ''' 4 | from .PerturbationObject import * 5 | from .SequentialAttacker import * 6 | from .ParallelAttacker import * 7 | from .CommonWeakness import * -------------------------------------------------------------------------------- /attacks/perturbation/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def cosine_similarity(x: list): 5 | ''' 6 | input a list of tensor with same shape. return the mean cosine_similarity 7 | ''' 8 | x = torch.stack(x, dim=0) 9 | N = x.shape[0] 10 | x = x.reshape(N, -1) 11 | 12 | norm = torch.norm(x, p=2, dim=1) 13 | x /= norm.reshape(-1, 1) # N, D 14 | similarity = x @ x.T # N, N 15 | mask = torch.triu(torch.ones(N, N, device=x.device), diagonal=0).to(torch.bool) 16 | similarity = similarity[mask] 17 | return torch.mean(similarity).item() 18 | 19 | def clamp(x: torch.tensor, min_value: float = 0, max_value: float = 1): 20 | return torch.clamp(x, min=min_value, max=max_value) 21 | -------------------------------------------------------------------------------- /attacks/stadv_eot/__init__.py: -------------------------------------------------------------------------------- 1 | from .attacks import StAdvAttack 2 | -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/LICENSE_RECOLORADV: -------------------------------------------------------------------------------- 1 | MIT License 2 | Copyright (c) 2018 YOUR NAME 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | The above copyright notice and this permission notice shall be included in all 10 | copies or substantial portions of the Software. 11 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 12 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 13 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 14 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 15 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 16 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 17 | SOFTWARE. -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/__init__.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from recoloradv. 5 | # 6 | # Source: 7 | # https://github.com/cassidylaidlaw/ReColorAdv/blob/master/recoloradv/__init__.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in this directory (LICENSE_RECOLORADV). 11 | # The modifications to this file are subject to the same license. 12 | # --------------------------------------------------------------- 13 | -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/mister_ed/README.md: -------------------------------------------------------------------------------- 1 | Code in this directory is adapted from the [`mister_ed`](https://github.com/revbucket/mister_ed) library. -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/mister_ed/__init__.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from recoloradv. 5 | # 6 | # Source: 7 | # https://github.com/cassidylaidlaw/ReColorAdv/blob/master/recoloradv/mister_ed/__init__.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in the `recoloradv` directory (LICENSE_RECOLORADV). 11 | # The modifications to this file are subject to the same license. 12 | # --------------------------------------------------------------- 13 | -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/mister_ed/config.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from recoloradv. 5 | # 6 | # Source: 7 | # https://github.com/cassidylaidlaw/ReColorAdv/blob/master/recoloradv/mister_ed/config.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in the `recoloradv` directory (LICENSE_RECOLORADV). 11 | # The modifications to this file are subject to the same license. 12 | # --------------------------------------------------------------- 13 | 14 | import os 15 | 16 | config_dir = os.path.abspath(os.path.dirname(__file__)) 17 | 18 | 19 | def path_resolver(path): 20 | if path.startswith('~/'): 21 | return os.path.expanduser(path) 22 | 23 | if path.startswith('./'): 24 | return os.path.join(*[config_dir] + path.split('/')[1:]) 25 | 26 | 27 | DEFAULT_DATASETS_DIR = path_resolver('~/datasets') 28 | MODEL_PATH = path_resolver('./pretrained_models/') 29 | OUTPUT_IMAGE_PATH = path_resolver('./output_images/') 30 | 31 | DEFAULT_BATCH_SIZE = 128 32 | DEFAULT_WORKERS = 4 33 | CIFAR10_MEANS = [0.485, 0.456, 0.406] 34 | CIFAR10_STDS = [0.229, 0.224, 0.225] 35 | 36 | WIDE_CIFAR10_MEANS = [0.4914, 0.4822, 0.4465] 37 | WIDE_CIFAR10_STDS = [0.2023, 0.1994, 0.2010] 38 | 39 | IMAGENET_MEANS = [0.485, 0.456, 0.406] 40 | IMAGENET_STDS = [0.229, 0.224, 0.225] 41 | -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/mister_ed/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from recoloradv. 5 | # 6 | # Source: 7 | # https://github.com/cassidylaidlaw/ReColorAdv/blob/master/recoloradv/mister_ed/utils/__init__.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in the `recoloradv` directory (LICENSE_RECOLORADV). 11 | # The modifications to this file are subject to the same license. 12 | # --------------------------------------------------------------- 13 | -------------------------------------------------------------------------------- /attacks/stadv_eot/recoloradv/norms.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from recoloradv. 5 | # 6 | # Source: 7 | # https://github.com/cassidylaidlaw/ReColorAdv/blob/master/recoloradv/norms.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in this directory (LICENSE_RECOLORADV). 11 | # The modifications to this file are subject to the same license. 12 | # --------------------------------------------------------------- 13 | 14 | import torch 15 | from torch.autograd import Variable 16 | 17 | 18 | def smoothness(grid): 19 | """ 20 | Given a variable of dimensions (N, X, Y, [Z], C), computes the sum of 21 | the differences between adjacent points in the grid formed by the 22 | dimensions X, Y, and (optionally) Z. Returns a tensor of dimension N. 23 | """ 24 | 25 | num_dims = len(grid.size()) - 2 26 | batch_size = grid.size()[0] 27 | norm = Variable(torch.zeros(batch_size, dtype=grid.data.dtype, 28 | device=grid.data.device)) 29 | 30 | for dim in range(num_dims): 31 | slice_before = (slice(None),) * (dim + 1) 32 | slice_after = (slice(None),) * (num_dims - dim) 33 | shifted_grids = [ 34 | # left 35 | torch.cat([ 36 | grid[slice_before + (slice(1, None),) + slice_after], 37 | grid[slice_before + (slice(-1, None),) + slice_after], 38 | ], dim + 1), 39 | # right 40 | torch.cat([ 41 | grid[slice_before + (slice(None, 1),) + slice_after], 42 | grid[slice_before + (slice(None, -1),) + slice_after], 43 | ], dim + 1) 44 | ] 45 | for shifted_grid in shifted_grids: 46 | delta = shifted_grid - grid 47 | norm_components = (delta.pow(2).sum(-1) + 1e-10).pow(0.5) 48 | norm.add_(norm_components.sum( 49 | tuple(range(1, len(norm_components.size()))))) 50 | 51 | return norm 52 | -------------------------------------------------------------------------------- /attacks/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | 4 | 5 | def clamp(x: torch.tensor, min_value=0, max_value=1): 6 | return torch.clamp(x, min=min_value, max=max_value) 7 | 8 | 9 | def inplace_clamp(x: torch.tensor, min_value: float = 0, max_value: float = 1): 10 | return x.clamp_(min=min_value, max=max_value) 11 | -------------------------------------------------------------------------------- /data/NIPS17.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset, DataLoader 2 | import csv 3 | import os 4 | from PIL import Image 5 | from torchvision import transforms 6 | 7 | __kaggle_link__ = 'kaggle datasets download -d google-brain/nips-2017-adversarial-learning-development-set' 8 | 9 | 10 | class NIPS17(Dataset): 11 | def __init__(self, images_path='./resources/NIPS17/images/', 12 | label_path='./resources/NIPS17/images.csv', 13 | transform=transforms.Compose([ 14 | transforms.Resize((224, 224)), 15 | transforms.ToTensor(), 16 | ]), 17 | ): 18 | self.labels = {} 19 | with open(label_path) as f: 20 | reader = csv.reader(f) 21 | for line in list(reader)[1:]: 22 | name, label = line[0], int(line[6]) - 1 23 | self.labels[name + '.png'] = label 24 | self.images = os.listdir(images_path) 25 | self.images.sort() 26 | self.images_path = images_path 27 | self.transform = transform 28 | 29 | def __len__(self): 30 | return len(self.images) 31 | 32 | def __getitem__(self, item): 33 | name = self.images[item] 34 | x = Image.open(os.path.join(self.images_path, name)) 35 | y = self.labels[name] 36 | return self.transform(x), y 37 | 38 | 39 | def get_NIPS17_loader(batch_size=64, 40 | num_workers=8, 41 | pin_memory=True, 42 | download=False, 43 | shuffle=False, 44 | transform=transforms.Compose([ 45 | transforms.Resize((224, 224)), 46 | transforms.ToTensor(), 47 | ]), 48 | **kwargs, 49 | ): 50 | if download: 51 | os.system(__kaggle_link__) 52 | set = NIPS17(transform=transform, **kwargs) 53 | loader = DataLoader(set, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, 54 | shuffle=shuffle) 55 | return loader 56 | -------------------------------------------------------------------------------- /data/PACS.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader, Dataset, ConcatDataset 3 | import torchvision 4 | from torchvision import transforms 5 | from tllib.vision.datasets import PACS 6 | from tllib.vision.transforms import ResizeImage 7 | from tllib.vision.datasets.imagelist import MultipleDomainsDataset 8 | 9 | """ 10 | install tllib: 11 | git clone git@github.com:thuml/Transfer-Learning-Library.git 12 | python setup.py install 13 | pip install -r requirements.txt 14 | """ 15 | 16 | 17 | class NPACS(PACS): 18 | def __init__(self, root: str, task: str, split='all', download=True, **kwargs): 19 | super(NPACS, self).__init__(root, task, split, download, **kwargs) 20 | 21 | def __getitem__(self, index): 22 | img, target = super(NPACS, self).__getitem__(index) 23 | return img, target, str(index) 24 | 25 | 26 | def get_pacs_dataset(target_domain, root="./data/pacs", download=True, augment=True): 27 | assert target_domain in ["P", "A", "C", "S"] 28 | test_transform = transforms.Compose( 29 | [ 30 | ResizeImage(224), 31 | transforms.ToTensor(), 32 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 33 | 34 | ] 35 | ) 36 | train_transform = transforms.Compose( 37 | [ 38 | transforms.RandomResizedCrop(224, scale=(0.7, 1.0)), 39 | # transforms.RandomHorizontalFlip(), 40 | # transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3), 41 | # transforms.RandomGrayscale(), 42 | transforms.AutoAugment(), 43 | transforms.ToTensor(), 44 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 45 | 46 | ] 47 | ) 48 | test_dataset = NPACS(root=root, 49 | task=target_domain, 50 | transform=test_transform, 51 | download=download) 52 | 53 | source_domain = [i for i in ["P", "A", "C", "S"] if target_domain != i] 54 | 55 | train_dataset = [] 56 | for domain in source_domain: 57 | train_dataset.append(NPACS(root=root, 58 | task=domain, 59 | transform=train_transform if augment else test_transform, 60 | download=download)) 61 | train_dataset = ConcatDataset(train_dataset) 62 | return train_dataset, test_dataset 63 | 64 | 65 | def get_PACS_train(batch_size=32, 66 | num_workers=8, 67 | pin_memory=True, 68 | augment=True, 69 | target_domain="A" 70 | ): 71 | set, _ = get_pacs_dataset(root='./resource/PACS', target_domain=target_domain, augment=augment) 72 | loader = DataLoader(set, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, 73 | shuffle=True) 74 | return loader 75 | 76 | 77 | def get_PACS_test(batch_size=32, 78 | num_workers=8, 79 | pin_memory=True, 80 | target_domain="A" 81 | ): 82 | _, set = get_pacs_dataset(root='./resource/PACS', target_domain=target_domain) 83 | loader = DataLoader(set, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, 84 | shuffle=True) 85 | return loader 86 | -------------------------------------------------------------------------------- /data/__init__.py: -------------------------------------------------------------------------------- 1 | from .cifar import * 2 | from .NIPS17 import * 3 | from .someset import get_someset_loader, save_to_someset 4 | from .CIFAR10C import get_cifar_10_c_loader 5 | from .ImageNet import get_imagenet10_loader, get_imagenet_loader, get_imagenet_subset 6 | from .mnist import * 7 | -------------------------------------------------------------------------------- /data/mnist.py: -------------------------------------------------------------------------------- 1 | from torchvision.datasets import MNIST 2 | from torchvision import transforms 3 | from torch.utils.data import DataLoader 4 | 5 | __all__ = ['get_mnist_test', 'get_mnist_train'] 6 | 7 | 8 | def get_mnist_train(batch_size=256, 9 | num_workers=8, 10 | pin_memory=True, 11 | ): 12 | transform = transforms.Compose([ 13 | transforms.Resize((32, 32)), 14 | transforms.ToTensor(), 15 | ]) 16 | set = MNIST('./resources/mnist/', train=True, download=True, transform=transform) 17 | loader = DataLoader(set, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory, 18 | shuffle=True) 19 | return loader 20 | 21 | 22 | def get_mnist_test(batch_size=256, 23 | num_workers=8, 24 | pin_memory=True, ): 25 | transform = transforms.Compose([ 26 | transforms.Resize((32, 32)), 27 | transforms.ToTensor(), 28 | ]) 29 | set = MNIST('./resources/mnist/', train=False, download=True, transform=transform) 30 | loader = DataLoader(set, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory) 31 | return loader 32 | -------------------------------------------------------------------------------- /data/someset.py: -------------------------------------------------------------------------------- 1 | """ 2 | this file aims to read any dataset satisfied that: 3 | 1.all the images are in one folder 4 | 2.only a dict to store ground truth. Keys are image names, values are ground truth labels. 5 | """ 6 | 7 | import os 8 | import torch 9 | from torch.utils.data import DataLoader, Dataset 10 | from torchvision import transforms 11 | from PIL import Image 12 | import numpy as np 13 | from torch import Tensor 14 | from typing import List 15 | from utils import save_list_images 16 | 17 | 18 | def sort_key(x: str) -> int: 19 | return int(x[:-4]) 20 | 21 | 22 | class SomeDataSet(Dataset): 23 | def __init__(self, img_path, gt_path, transform=transforms.ToTensor()): 24 | self.transform = transform 25 | self.images = [img for img in os.listdir(img_path) if img.endswith(".jpg") or img.endswith(".png")] 26 | self.images.sort(key=sort_key) 27 | self.gt = np.load(gt_path, allow_pickle=True).item() 28 | self.img_path = img_path 29 | 30 | def __len__(self): 31 | return len(self.images) 32 | 33 | def __getitem__(self, item): 34 | now = self.images[item] 35 | now_img = Image.open(os.path.join(self.img_path, now)) # numpy 36 | return self.transform(now_img), self.gt[now] 37 | 38 | 39 | def get_someset_loader( 40 | img_path, 41 | gt_path, 42 | batch_size=256, 43 | num_workers=8, 44 | pin_memory=False, 45 | shuffle=False, 46 | transform=transforms.ToTensor(), 47 | ): 48 | set = SomeDataSet(img_path=img_path, gt_path=gt_path, transform=transform) 49 | loader = DataLoader(set, batch_size=batch_size, num_workers=num_workers, shuffle=shuffle, pin_memory=pin_memory) 50 | return loader 51 | 52 | 53 | def save_to_someset(images: List[Tensor], gts: List[int], path="./debug/someset"): 54 | save_list_images(images, path) 55 | labels = dict() 56 | for i, gt in enumerate(gts): 57 | labels[f"{i}.png"] = gt 58 | np.save(os.path.join(path, "labels.npy"), labels) 59 | -------------------------------------------------------------------------------- /data/utils.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset, DataLoader 2 | from torch import Tensor 3 | from utils.ImageHandling import save_list_images 4 | import numpy as np 5 | import os 6 | 7 | 8 | def get_loader(dataset: Dataset, 9 | batch_size=32, 10 | shuffle=True, 11 | ): 12 | loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle) 13 | return loader 14 | 15 | 16 | def save_dataset( 17 | x: Tensor, 18 | y: Tensor, 19 | path: str, 20 | gt_saving_name: str, 21 | ): 22 | if not os.path.exists(path): 23 | os.makedirs(path) 24 | ground_truth = dict() 25 | save_list_images(x.split(1, dim=0), folder_path=path) 26 | for i in range(y.shape[0]): 27 | ground_truth[str(i) + '.png'] = y[i].item() 28 | np.save(os.path.join(path, gt_saving_name), ground_truth) 29 | print(f'Successfully create a new dataset with {len(ground_truth)} images') 30 | -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/DYPUnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from typing import Tuple 4 | 5 | 6 | class BasicDownBlock(nn.Module): 7 | def __init__(self, in_dim: int, out_dim: int, stride: int = 1, padding: int = 1): 8 | super(BasicDownBlock, self).__init__() 9 | self.m = nn.Sequential( 10 | nn.Conv2d(in_dim, out_dim, kernel_size=3, stride=stride, padding=padding), 11 | nn.BatchNorm2d(out_dim), 12 | nn.ReLU(), 13 | ) 14 | 15 | def forward(self, x): 16 | return self.m(x) 17 | 18 | 19 | class C2(nn.Module): 20 | def __init__(self, in_dim: int, out_dim: int): 21 | super(C2, self).__init__() 22 | middle: int = (in_dim + out_dim) // 2 23 | self.m1 = BasicDownBlock(in_dim, middle) 24 | self.m2 = BasicDownBlock(middle, out_dim) 25 | 26 | def forward(self, x): 27 | x = self.m1(x) 28 | x = self.m2(x) 29 | return x 30 | 31 | 32 | class C3(nn.Module): 33 | def __init__(self, in_dim: int, out_dim: int, down=True): 34 | super(C3, self).__init__() 35 | if down: 36 | self.down = BasicDownBlock(in_dim, out_dim, stride=2) 37 | else: 38 | self.down = BasicDownBlock(in_dim, out_dim, stride=1) 39 | self.fuse = C2(out_dim, out_dim) 40 | 41 | def forward(self, x): 42 | x = self.down(x) 43 | x = self.fuse(x) 44 | return x 45 | 46 | 47 | class Fuse(nn.Module): 48 | def __init__(self, size: Tuple[int, int], in_dim, out_dim): 49 | super(Fuse, self).__init__() 50 | self.up = nn.Upsample(size=size, mode='bilinear') 51 | self.c3 = C3(in_dim, out_dim, down=False) 52 | 53 | def forward(self, x, res): 54 | x = self.up(x) 55 | x = torch.cat([x, res], dim=1) 56 | x = self.c3(x) 57 | return x 58 | 59 | 60 | class DYPUnet(nn.Module): 61 | def __init__(self): 62 | super(DYPUnet, self).__init__() 63 | self.first = C2(3, 64) 64 | self.second = C3(64, 128) 65 | self.third = C3(128, 256) 66 | self.fourth = C3(256, 256) 67 | self.fifth = C3(256, 256) 68 | self.f1 = Fuse((38, 38), 512, 256) 69 | self.f2 = Fuse((75, 75), 512, 256) 70 | self.f3 = Fuse((150, 150), 384, 128) 71 | self.f4 = Fuse((299, 299), 192, 64) 72 | self.final = nn.Conv2d(64, 3, stride=1, kernel_size=1) 73 | 74 | def forward(self, x): 75 | x1 = self.first(x) 76 | x2 = self.second(x1) 77 | x3 = self.third(x2) 78 | x4 = self.fourth(x3) 79 | x5 = self.fifth(x4) 80 | # print(x5.shape) 81 | x4 = self.f1(x5, x4) 82 | x3 = self.f2(x4, x3) 83 | x2 = self.f3(x3, x2) 84 | x1 = self.f4(x2, x1) 85 | x1 = self.final(x1) # - adv perturb 86 | return x + x1 87 | 88 | 89 | if __name__ == '__main__': 90 | idol = DYPUnet() 91 | x = torch.randn(1, 3, 299, 299) 92 | print(x.shape) 93 | print(idol(x).shape) -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | this module aims to reproduce 3 | Defense against Adversarial Attacks Using High-Level Representation Guided HighLevelGuidedDenoiser 4 | in pytorch. 5 | ''' 6 | from .get_HGD import get_HGD_model 7 | -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/dataset.py: -------------------------------------------------------------------------------- 1 | import torch.utils.data as data 2 | 3 | from PIL import Image 4 | import os 5 | import re 6 | import torch 7 | 8 | IMG_EXTENSIONS = ['.png', '.jpg', '.jpeg'] 9 | 10 | 11 | def natural_key(string_): 12 | """See http://www.codinghorror.com/blog/archives/001018.html""" 13 | return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())] 14 | 15 | 16 | def find_images_and_targets(folder, types=IMG_EXTENSIONS, class_to_idx=None, leaf_name_only=True, sort=True): 17 | if class_to_idx is None: 18 | class_to_idx = dict() 19 | build_class_idx = True 20 | else: 21 | build_class_idx = False 22 | labels = [] 23 | filenames = [] 24 | for root, subdirs, files in os.walk(folder, topdown=False): 25 | rel_path = os.path.relpath(root, folder) if (root != folder) else '' 26 | label = os.path.basename(rel_path) if leaf_name_only else rel_path.replace(os.path.sep, '_') 27 | if build_class_idx and not subdirs: 28 | class_to_idx[label] = 0 29 | for f in files: 30 | base, ext = os.path.splitext(f) 31 | if ext.lower() in types: 32 | filenames.append(os.path.join(root,f)) 33 | labels.append(label) 34 | if build_class_idx: 35 | classes = sorted(class_to_idx.keys(), key=natural_key) 36 | for idx, c in enumerate(classes): 37 | class_to_idx[c] = idx 38 | images_and_targets = zip(filenames, [class_to_idx[l] for l in labels]) 39 | if sort: 40 | images_and_targets = sorted(images_and_targets, key=lambda k: natural_key(k[0])) 41 | if build_class_idx: 42 | return images_and_targets, classes, class_to_idx 43 | else: 44 | return images_and_targets 45 | 46 | 47 | class Dataset(data.Dataset): 48 | 49 | def __init__(self, root, transform=None): 50 | imgs, classes, class_to_idx = find_images_and_targets(root) 51 | if len(imgs) == 0: 52 | raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n" 53 | "Supported image extensions are: " + ",".join(IMG_EXTENSIONS))) 54 | 55 | self.root = root 56 | self.imgs = imgs 57 | self.classes = classes 58 | self.class_to_idx = class_to_idx 59 | self.transform = transform 60 | 61 | def __getitem__(self, index): 62 | path, target = self.imgs[index] 63 | img = Image.open(path).convert('RGB') 64 | if self.transform is not None: 65 | img = self.transform(img) 66 | if target is None: 67 | target = torch.zeros(1).long() 68 | return img, target 69 | 70 | def __len__(self): 71 | return len(self.imgs) 72 | 73 | def filenames(self, indices=[]): 74 | if indices: 75 | return [self.imgs[i][0] for i in indices] 76 | else: 77 | return [x[0] for x in self.imgs] 78 | -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/inres.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from .inceptionresnet import get_net, Conv, Bottleneck 4 | 5 | config = dict() 6 | config["flip"] = True 7 | config["loss_idcs"] = [4] 8 | config["net_type"] = "inceptionresnetv2" 9 | 10 | input_size = [299, 299] 11 | block = [Conv, None, None, None, None] 12 | fwd_out = [[64, 128, 256, 256, 256], [], [], [], []] 13 | num_fwd = [[2, 3, 3, 3, 3], [], [], [], []] 14 | back_out = [[64, 128, 256, 256], [], [], [], []] 15 | num_back = [[2, 3, 3, 3], [], [], [], []] 16 | n = 1 17 | hard_mining = 0 18 | loss_norm = False 19 | 20 | 21 | def get_model(): 22 | net = get_net(input_size, block, fwd_out, num_fwd, back_out, num_back, n, hard_mining, loss_norm) 23 | return config, net 24 | -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/res152_wide.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from .resnet import get_net, Conv, Bottleneck 4 | 5 | config = dict() 6 | config["flip"] = True 7 | config["loss_idcs"] = [4] 8 | 9 | net_type = "resnet152" 10 | input_size = [299, 299] 11 | block = [Conv, None, None, None, None] 12 | fwd_out = [[64, 128, 256, 256, 256], [128, 256, 256, 256], [], [], []] 13 | num_fwd = [[2, 3, 3, 3, 3], [2, 3, 3, 3], [], [], []] 14 | back_out = [[64, 128, 256, 256], [128, 256, 256], [], [], []] 15 | num_back = [[2, 3, 3, 3], [2, 3, 3], [], [], []] 16 | n = 1 17 | hard_mining = 0 18 | loss_norm = False 19 | 20 | 21 | def get_model(): 22 | net = get_net(net_type, input_size, block, fwd_out, num_fwd, back_out, num_back, n, hard_mining, loss_norm) 23 | return config, net 24 | -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/resnext101.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from .resnext import get_net, Conv, Bottleneck 4 | 5 | config = dict() 6 | config['flip'] = True 7 | config['loss_idcs'] = [1] 8 | 9 | net_type = 'resnext101' 10 | config['net_type'] = net_type 11 | input_size = [299, 299] 12 | block = Conv 13 | fwd_out = [64, 128, 256, 256, 256] 14 | num_fwd = [2, 3, 3, 3, 3] 15 | back_out = [64, 128, 256, 256] 16 | num_back = [2, 3, 3, 3] 17 | n = 1 18 | hard_mining = 0 19 | loss_norm = False 20 | 21 | def get_model(): 22 | net = get_net(input_size, block, fwd_out, num_fwd, back_out, num_back, n, hard_mining, loss_norm) 23 | return config, net 24 | -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/resnext_features/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnext101_32x4d_features import resnext101_32x4d_features 2 | from .resnext101_64x4d_features import resnext101_64x4d_features -------------------------------------------------------------------------------- /defenses/HighLevelGuidedDenoiser/v3.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from .inception import get_net, Conv, Bottleneck 4 | 5 | config = dict() 6 | config["flip"] = True 7 | config["loss_idcs"] = [1] 8 | config["net_type"] = "inceptionresnetv2" 9 | input_size = [299, 299] 10 | block = Conv 11 | fwd_out = [64, 128, 256, 256, 256] 12 | num_fwd = [2, 3, 3, 3, 3] 13 | back_out = [64, 128, 256, 256] 14 | num_back = [2, 3, 3, 3] 15 | n = 1 16 | hard_mining = 0 17 | loss_norm = False 18 | 19 | 20 | def get_model(): 21 | net = get_net(input_size, block, fwd_out, num_fwd, back_out, num_back, n, hard_mining, loss_norm) 22 | return config, net 23 | -------------------------------------------------------------------------------- /defenses/NeuralRepresentationPurifier/NRP.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .networks import NRP, NRP_resG 3 | from torchvision import transforms 4 | 5 | 6 | class NeuralRepresentationPurifier(torch.nn.Module): 7 | def __init__(self, model: torch.nn.Module, 8 | transform=transforms.Compose([ 9 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 10 | ]), 11 | mode='NRP_resG', 12 | ): 13 | super(NeuralRepresentationPurifier, self).__init__() 14 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 15 | self.model = model 16 | self.transforms = transform 17 | if mode == 'NRP': 18 | self.netG = NRP(3, 3, 64, 23) 19 | self.netG.load_state_dict(torch.load('resources/checkpoints/NRP/NRP.pth')) 20 | if mode == 'NRP_resG': 21 | self.netG = NRP_resG(3, 3, 64, 23) 22 | self.netG.load_state_dict(torch.load('resources/checkpoints/NRP/NRP_resG.pth')) 23 | self.netG.to(device=self.device) 24 | self.i = 0 25 | self.to_img = transforms.ToPILImage() 26 | 27 | def forward(self, x): 28 | x = self.netG(x) 29 | return self.model(x) 30 | 31 | -------------------------------------------------------------------------------- /defenses/NeuralRepresentationPurifier/__init__.py: -------------------------------------------------------------------------------- 1 | from .NRP import NeuralRepresentationPurifier 2 | -------------------------------------------------------------------------------- /defenses/NeuralRepresentationPurifier/modules/module_util.py: -------------------------------------------------------------------------------- 1 | # Courtsy of https://github.com/xinntao/BasicSR 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.init as init 5 | import torch.nn.functional as F 6 | 7 | 8 | def initialize_weights(net_l, scale=1): 9 | if not isinstance(net_l, list): 10 | net_l = [net_l] 11 | for net in net_l: 12 | for m in net.modules(): 13 | if isinstance(m, nn.Conv2d): 14 | init.kaiming_normal_(m.weight, a=0, mode='fan_in') 15 | m.weight.data *= scale # for residual block 16 | if m.bias is not None: 17 | m.bias.data.zero_() 18 | elif isinstance(m, nn.Linear): 19 | init.kaiming_normal_(m.weight, a=0, mode='fan_in') 20 | m.weight.data *= scale 21 | if m.bias is not None: 22 | m.bias.data.zero_() 23 | elif isinstance(m, nn.BatchNorm2d): 24 | init.constant_(m.weight, 1) 25 | init.constant_(m.bias.data, 0.0) 26 | 27 | 28 | def make_layer(block, n_layers): 29 | layers = [] 30 | for _ in range(n_layers): 31 | layers.append(block()) 32 | return nn.Sequential(*layers) 33 | 34 | 35 | class ResidualBlock_noBN(nn.Module): 36 | '''Residual block w/o BN 37 | ---Conv-ReLU-Conv-+- 38 | |________________| 39 | ''' 40 | 41 | def __init__(self, nf=64): 42 | super(ResidualBlock_noBN, self).__init__() 43 | self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) 44 | self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) 45 | 46 | # initialization 47 | initialize_weights([self.conv1, self.conv2], 0.1) 48 | 49 | def forward(self, x): 50 | identity = x 51 | out = F.relu(self.conv1(x), inplace=True) 52 | out = self.conv2(out) 53 | return identity + out 54 | 55 | 56 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffPure/DiffusionLikelihoodMaximizer.py: -------------------------------------------------------------------------------- 1 | from torch import nn, Tensor 2 | from torch.autograd import Function 3 | 4 | from defenses.PurificationDefenses.PurificationDefense import PurificationDefense 5 | from typing import Callable 6 | 7 | 8 | class DiffusionLikelihoodMaximizerFunction(Function): 9 | function = None 10 | 11 | @staticmethod 12 | def forward(ctx, x: Tensor): 13 | x = DiffusionLikelihoodMaximizerFunction.function(x.detach().clone()) 14 | return x 15 | 16 | @staticmethod 17 | def backward(ctx, grad): 18 | return grad 19 | 20 | 21 | def diffusion_likelihood_maximizer_defense( 22 | classifier: nn.Module, 23 | likelihood_maximization: Callable, 24 | ): 25 | DiffusionLikelihoodMaximizerFunction.function = likelihood_maximization 26 | 27 | class Purifier(nn.Module): 28 | def __init__(self): 29 | super(Purifier, self).__init__() 30 | self.requires_grad_(False).eval().cuda() 31 | 32 | def forward(self, x): 33 | return DiffusionLikelihoodMaximizerFunction.apply(x) 34 | 35 | return PurificationDefense(Purifier(), classifier) 36 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffPure/RS.py: -------------------------------------------------------------------------------- 1 | from .diffpure import DiffusionPure 2 | import torch 3 | 4 | 5 | class DiffPureForRandomizedSmoothing(DiffusionPure): 6 | def __init__(self, *args, sigma=0.25, **kwargs): 7 | """ 8 | explanation for computing t: 9 | Sigma need to be multiplied by 2 because RS add noise at [0, 1] while diffusion working at [-1, 1] 10 | Because alpha_t = 1 / (sigma ** 2 + 1), we can compute t by: 11 | alpha[t] >= 1 / (sigma ** 2 + 1) > alpha[t+1] 12 | """ 13 | super(DiffPureForRandomizedSmoothing, self).__init__(*args, **kwargs) 14 | self.t = torch.max(self.diffusion.alpha_bar < 1 / ((sigma * 2) ** 2 + 1), dim=0)[1] - 1 15 | print(self.t) 16 | 17 | def forward(self, x, *args, **kwargs): 18 | x = self.pre_transforms(x) 19 | x = self.diffusion(x, *args, noise_level=(self.t + 1), scale=True, add_noise=False, **kwargs) 20 | x = self.post_transforms(x) 21 | x = self.model(x) 22 | return x 23 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffPure/__init__.py: -------------------------------------------------------------------------------- 1 | from .diffpure import DiffusionPure 2 | from .DiffusionLikelihoodMaximizer import diffusion_likelihood_maximizer_defense 3 | from .RS import DiffPureForRandomizedSmoothing 4 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffPure/diffpure.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from models import WideResNet_70_16_dropout 4 | from ..sampler import DDIM, DiffusionOde, DiffusionSde, EDMStochasticSampler, BaseSampler 5 | 6 | 7 | class DiffusionPure(nn.Module): 8 | def __init__(self, mode='sde', 9 | sampler: BaseSampler = None, 10 | pre_transforms=nn.Identity(), 11 | post_transforms=nn.Identity(), 12 | model=None, 13 | *args, **kwargs): 14 | super(DiffusionPure, self).__init__() 15 | self.device = torch.device('cuda') 16 | if sampler is not None: 17 | self.diffusion = sampler 18 | else: 19 | if mode == 'sde': 20 | self.diffusion = DiffusionSde(*args, **kwargs) 21 | elif mode == 'ode': 22 | self.diffusion = DiffusionOde(*args, **kwargs) 23 | elif mode == 'ddim': 24 | self.diffusion = DDIM(*args, **kwargs) 25 | elif mode == 'edm': 26 | self.diffusion = EDMStochasticSampler(*args, **kwargs) 27 | elif mode == 'ddpm': 28 | self.diffusion = DDIM(*args, ddpm=True, **kwargs) 29 | if model is None: 30 | model = WideResNet_70_16_dropout() 31 | self.model = model 32 | self.pre_transforms = pre_transforms 33 | self.post_transforms = post_transforms 34 | self.init() 35 | 36 | def init(self): 37 | self.eval().requires_grad_(False) 38 | self.to(self.device) 39 | 40 | def forward(self, x, *args, **kwargs): 41 | x = self.pre_transforms(x) 42 | x = self.diffusion(x, *args, **kwargs) 43 | x = self.post_transforms(x) 44 | x = self.model(x) 45 | return x 46 | 47 | 48 | class SerialDiffusionPure(DiffusionPure): 49 | def __init__(self, *args, **kwargs): 50 | self.t = kwargs['diffusion_number'] 51 | del kwargs['diffusion_number'] 52 | kwargs['dt'] = 1e-3 53 | super(SerialDiffusionPure, self).__init__(*args, **kwargs) 54 | 55 | def forward(self, x, *args, **kwargs): 56 | x = self.pre_transforms(x) 57 | for _ in range(self.t): 58 | x = self.diffusion(x, *args, **kwargs) 59 | x = self.post_transforms(x) 60 | x = self.model(x) 61 | return x 62 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffusionClassifier/OptimalDiffusionClassifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | 4 | 5 | class OptimalDiffusionClassifier(nn.Module): 6 | def __init__(self, 7 | loader, 8 | beta=torch.linspace(0.1 / 1000, 20 / 1000, 1000), 9 | num_classes=10, 10 | device=torch.device('cuda')): 11 | super(OptimalDiffusionClassifier, self).__init__() 12 | self.device = device 13 | self.num_classes = num_classes 14 | alpha = (1 - beta) 15 | self.T = alpha.numel() 16 | self.alpha_bar = alpha.cumprod(dim=0).to(self.device) 17 | self.sigma = torch.sqrt(1 - self.alpha_bar) 18 | self.loader = loader 19 | self._init() 20 | 21 | def _init(self): 22 | xs, ys = [], [] 23 | for x, y in self.loader: 24 | xs.append(x.to(self.device)) 25 | ys.append(y.to(self.device)) 26 | self.xs, self.ys = torch.cat(xs, dim=0), torch.cat(ys, dim=0) 27 | print('Finish initialize optimal diffusion classifier!!!') 28 | print('-' * 100) 29 | 30 | def forward_one_logit(self, x, y: int) -> Tensor: 31 | """ 32 | Here, we still do not need to Monte Carlo epsilon 33 | :param x: 1, C, H, D 34 | :param y: 1 35 | :return: 36 | """ 37 | mask = self.ys == y 38 | xs = self.xs[mask] # N, C, H, D 39 | delta = x - xs # N, C, H, D 40 | 41 | repeated_x = x.repeat(self.T, 1, 1, 1) 42 | xts = torch.sqrt(self.alpha_bar).view(-1, 1, 1, 1) * repeated_x + \ 43 | self.sigma.view(-1, 1, 1, 1) * torch.randn_like(repeated_x) # T, C, H, D 44 | softmax_inner = xts - torch.sqrt(self.alpha_bar).view(1, -1, 1, 1, 1) \ 45 | * xs.unsqueeze(1).repeat(1, self.T, 1, 1, 1) # N, T, C, H, D 46 | N, T, C, H, D = softmax_inner.shape 47 | softmax_inner = torch.sum(softmax_inner.view(N, T, -1) ** 2, dim=2) # N, T 48 | softmaxs = torch.softmax(- softmax_inner / (2 * self.sigma.view(1, -1)), dim=0) # N, T 49 | softmaxs = softmaxs.permute(1, 0) # T, N 50 | 51 | delta = delta.view(1, N, C * H * D) 52 | inner = softmaxs.view(T, N, 1) * delta # T, N, C*H*D 53 | inner = torch.sum(inner.reshape(T, N * C * H * D) ** 2, dim=1) # T 54 | result = torch.mean(inner * self.alpha_bar / self.sigma ** 2) 55 | return result 56 | 57 | def get_one_instance_prediction(self, x) -> Tensor: 58 | logit = [] 59 | for i in range(self.num_classes): 60 | logit.append(self.forward_one_logit(x, i)) 61 | logit = torch.stack(logit) 62 | # print(logit) 63 | return logit * -1 64 | 65 | def forward(self, x) -> Tensor: 66 | xs = x.split(1) # 1, C, H, D 67 | y = [] 68 | for now_x in xs: 69 | y.append(self.get_one_instance_prediction(now_x)) 70 | y = torch.stack(y) # N, num_classes 71 | return y 72 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffusionClassifier/__init__.py: -------------------------------------------------------------------------------- 1 | from .DiffusionClassifier import DiffusionClassifier, RobustDiffusionClassifier 2 | from .DiffusionClassifierImageNet import DiffusionAsClassifierImageNetWraped 3 | from .OptimalDiffusionClassifier import OptimalDiffusionClassifier 4 | from .SBGC import SBGC 5 | from .EDMDC import EDMEulerIntegralDC, EDMEulerIntegralWraped 6 | from .DiffusionClassifierBase import DiffusionClassifierSingleHeadBaseWraped 7 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/DiffusionClassifier/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from typing import List 4 | 5 | 6 | def clamp(x, min=0, max=1): 7 | return torch.clamp(x, min=min, max=max) 8 | 9 | 10 | def inplace_clamp(x, min=0, max=1): 11 | return torch.clamp(x, min=min, max=max) 12 | 13 | 14 | def L2Loss(x: Tensor, y: Tensor) -> Tensor: 15 | """ 16 | :param x: N, C, H, D 17 | :param y: N, C, H, D 18 | :return: dim=0 tensor 19 | """ 20 | x = (x - y) ** 2 21 | x = x.view(x.shape[0], -1) 22 | x = torch.norm(x, dim=1, p=2) 23 | x = x.mean(0) 24 | return x 25 | 26 | 27 | def abs_loss(x: Tensor, y: Tensor) -> Tensor: 28 | diff = torch.abs(x - y) 29 | diff = diff.view(diff.shape[0], -1) 30 | diff = torch.sum(diff, dim=1) 31 | diff = torch.mean(diff, dim=0) 32 | return diff 33 | 34 | 35 | def L2_each_instance(x: Tensor, y: Tensor) -> Tensor: 36 | """ 37 | N, ?, ?, ?, ... 38 | :return: N, 39 | """ 40 | x = (x - y) ** 2 41 | x = x.view(x.shape[0], -1) 42 | x = torch.norm(x, dim=1, p=2) 43 | return x 44 | 45 | 46 | def list_mean(x: List) -> float: 47 | return sum(x) / len(x) 48 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/LikelihoodMaximization/LikelihoodMaximizationBase.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | import torch 3 | 4 | 5 | class LikelihoodMaximizationBase(nn.Module): 6 | def __init__(self, dc: nn.Module, device=torch.device("cuda"), transform=lambda x: (x - 0.5) * 2): 7 | super().__init__() 8 | self.dc = dc 9 | self.device = device 10 | self.transform = transform 11 | self.eval().requires_grad_(False).to(self.device) 12 | 13 | def optimize_back(self, *args, **kwargs): 14 | pass 15 | 16 | def forward(self, *args, **kwargs): 17 | return self.optimize_back(*args, **kwargs) 18 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/LikelihoodMaximization/__init__.py: -------------------------------------------------------------------------------- 1 | from .EDMLM import EDMEulerIntegralLM 2 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/LikelihoodMaximization/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from typing import List 4 | 5 | 6 | def clamp(x, min=0, max=1): 7 | return torch.clamp(x, min=min, max=max) 8 | 9 | 10 | def inplace_clamp(x, min=0, max=1): 11 | return torch.clamp(x, min=min, max=max) 12 | 13 | 14 | def L2Loss(x: Tensor, y: Tensor) -> Tensor: 15 | """ 16 | :param x: N, C, H, D 17 | :param y: N, C, H, D 18 | :return: dim=0 tensor 19 | """ 20 | x = (x - y) ** 2 21 | x = x.view(x.shape[0], -1) 22 | x = torch.norm(x, dim=1, p=2) 23 | x = x.mean(0) 24 | return x 25 | 26 | 27 | def abs_loss(x: Tensor, y: Tensor) -> Tensor: 28 | diff = torch.abs(x - y) 29 | diff = diff.view(diff.shape[0], -1) 30 | diff = torch.sum(diff, dim=1) 31 | diff = torch.mean(diff, dim=0) 32 | return diff 33 | 34 | 35 | def L2_each_instance(x: Tensor, y: Tensor) -> Tensor: 36 | """ 37 | N, ?, ?, ?, ... 38 | :return: N, 39 | """ 40 | x = (x - y) ** 2 41 | x = x.view(x.shape[0], -1) 42 | x = torch.norm(x, dim=1, p=2) 43 | return x 44 | 45 | 46 | def list_mean(x: List) -> float: 47 | return sum(x) / len(x) 48 | 49 | 50 | def l2_clamp(x: Tensor, ori_x: Tensor, eps=1.0, inplace=False) -> Tensor: 51 | """ 52 | Here, x should be in [-1, 1]. eps=1 actually corresponding to 0.5 in [0, 1] 53 | """ 54 | B = x.shape[0] 55 | difference = x - ori_x 56 | distance = torch.norm(difference.view(B, -1), p=2, dim=1) 57 | mask = distance > eps 58 | if torch.sum(mask) > 0: 59 | difference[mask] = difference[mask] / distance[mask].view(torch.sum(mask), 1, 1, 1) * eps 60 | if inplace: 61 | x.mul_(0).add_(ori_x + difference) 62 | else: 63 | x = ori_x + difference 64 | return x 65 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/__init__.py: -------------------------------------------------------------------------------- 1 | from .DiffPure import DiffusionPure, DiffPureForRandomizedSmoothing, diffusion_likelihood_maximizer_defense 2 | from .DiffusionClassifier import * 3 | from .sampler import * 4 | from .LikelihoodMaximization import * 5 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from .basic import * 2 | from .EDMLosses import * 3 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/model.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | import torch 3 | from torch import nn, Tensor 4 | from models.unets import get_NCSNPP 5 | from models.unets import get_guided_diffusion_unet 6 | 7 | 8 | def get_unet(mode='cifar') -> Tuple[nn.Module, Tensor, Tuple[int, int, int]]: 9 | if mode == 'imagenet': 10 | img_shape = (3, 256, 256) 11 | model = get_guided_diffusion_unet() 12 | model.load_state_dict(torch.load(f'./resources/checkpoints/DiffPure/256x256_diffusion_uncond.pt')) 13 | elif mode == 'cifar': 14 | img_shape = (3, 32, 32) 15 | model = get_NCSNPP() 16 | model.load_state_dict(torch.load('./resources/checkpoints/DiffPure/32x32_diffusion.pth'), strict=False) 17 | else: 18 | assert False, 'We only provide CIFAR10 and ImageNet diffusion unet model' 19 | betas = torch.linspace(0.1 / 1000, 20 / 1000, 1000, device=torch.device('cuda')) 20 | return model, betas, img_shape 21 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/sampler/BaseSampler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | from abc import abstractmethod 4 | from ..model import get_unet 5 | from torch.utils.checkpoint import checkpoint 6 | 7 | 8 | def define_checkpointed_unet(unet): 9 | class SubstituteUnet(torch.nn.Module): 10 | def __init__(self): 11 | super(SubstituteUnet, self).__init__() 12 | for attr in list(dir(unet)): 13 | if not attr.startswith('__') and not attr.startswith('_') and attr != 'forward': 14 | setattr(self, attr, getattr(unet, attr)) 15 | self.unet = unet 16 | 17 | def forward(self, *args, **kwargs): 18 | x = checkpoint(self.unet, *args, **kwargs) 19 | return x 20 | 21 | return SubstituteUnet() 22 | 23 | 24 | class BaseSampler(nn.Module): 25 | def __init__(self, unet: nn.Module = None, 26 | device=torch.device('cuda'), 27 | grad_checkpoint=False): 28 | super(BaseSampler, self).__init__() 29 | unet = unet if unet is not None else get_unet()[0] 30 | self.unet = define_checkpointed_unet(unet) if grad_checkpoint else unet 31 | self.device = device 32 | self._model_init() 33 | 34 | def _model_init(self): 35 | self.eval().requires_grad_(False).to(self.device) 36 | 37 | @abstractmethod 38 | def sample(self, *args, **kwargs): 39 | pass 40 | 41 | @abstractmethod 42 | def purify(self, *args, **kwargs): 43 | pass 44 | 45 | def __call__(self, *args, **kwargs): 46 | return self.purify(*args, **kwargs) 47 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/sampler/PreconditionChanger.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | 4 | 5 | class EDM2VP(nn.Module): 6 | def __init__( 7 | self, 8 | edm_unet: nn.Module, 9 | beta=torch.linspace(0.1 / 1000, 20 / 1000, 1000, device=torch.device("cuda")), 10 | device=torch.device("cuda"), 11 | ): 12 | self.device = device 13 | super(EDM2VP, self).__init__() 14 | self.unet = edm_unet 15 | alpha = 1 - beta 16 | self.alpha_bar = alpha.cumprod(dim=0).to(self.device) 17 | 18 | def forward(self, x, t, *args, **kwargs): 19 | """ 20 | VP Inputs 21 | :param x: 22 | :param t: 23 | :param args: 24 | :param kwargs: 25 | :return: VP predictions 26 | """ 27 | t = t.long() 28 | x0 = self.unet( 29 | x / torch.sqrt(self.alpha_bar[t]).view(-1, 1, 1, 1), 30 | torch.sqrt(1 - self.alpha_bar[t]) / torch.sqrt(self.alpha_bar[t]), 31 | *args, 32 | **kwargs 33 | ) 34 | epsilon = (x - torch.sqrt(self.alpha_bar[t]).view(-1, 1, 1, 1) * x0) / torch.sqrt(1 - self.alpha_bar[t]).view( 35 | -1, 1, 1, 1 36 | ) 37 | return epsilon 38 | 39 | 40 | class VP2EDM(nn.Module): 41 | def __init__( 42 | self, 43 | vp_unet: nn.Module, 44 | beta=torch.linspace(0.1 / 1000, 20 / 1000, 1000, device=torch.device("cuda")), 45 | device=torch.device("cuda"), 46 | ): 47 | self.device = device 48 | super(VP2EDM, self).__init__() 49 | self.unet = vp_unet 50 | alpha = 1 - beta 51 | self.alpha_bar = alpha.cumprod(dim=0).to(self.device) 52 | self.edm_sigma = torch.sqrt((1 - self.alpha_bar) / self.alpha_bar) 53 | 54 | def round_sigma(self, sigma): 55 | return torch.as_tensor(sigma) 56 | 57 | def forward(self, x, sigma, *args, **kwargs): 58 | """ 59 | EDM Inputs 60 | :param x: N, C, H, D 61 | :param sigma: N, 62 | :param args: 63 | :param kwargs: 64 | :return: EDM predictions 65 | """ 66 | # chose the first t greater than sigma 67 | t = torch.max((sigma.view(-1, 1) <= self.edm_sigma.repeat(sigma.numel(), 1)), dim=1)[1] 68 | vp_x_t = x * torch.sqrt(self.alpha_bar[t]).view(-1, 1, 1, 1) 69 | epsilon = self.unet(vp_x_t, t, *args, **kwargs) 70 | # x_0 = x - sigma * epsilon. Will this be better? 71 | x_0 = (vp_x_t - torch.sqrt(1 - self.alpha_bar[t]).view(-1, 1, 1, 1) * epsilon) / torch.sqrt( 72 | self.alpha_bar[t] 73 | ).view(-1, 1, 1, 1) 74 | return x_0 75 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/sampler/__init__.py: -------------------------------------------------------------------------------- 1 | from .sde import DiffusionSde, DiffusionOde 2 | from .ddim import DDIM 3 | from .edm import EDMStochasticSampler 4 | from .PreconditionChanger import * 5 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/DiffPure/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from typing import List 4 | 5 | 6 | def clamp(x, min=0, max=1): 7 | return torch.clamp(x, min=min, max=max) 8 | 9 | 10 | def inplace_clamp(x, min=0, max=1): 11 | return torch.clamp(x, min=min, max=max) 12 | 13 | 14 | def L2Loss(x: Tensor, y: Tensor) -> Tensor: 15 | """ 16 | :param x: N, C, H, D 17 | :param y: N, C, H, D 18 | :return: dim=0 tensor 19 | """ 20 | x = (x - y) ** 2 21 | x = x.view(x.shape[0], -1) 22 | x = torch.norm(x, dim=1, p=2) 23 | x = x.mean(0) 24 | return x 25 | 26 | 27 | def abs_loss(x: Tensor, y: Tensor) -> Tensor: 28 | diff = torch.abs(x - y) 29 | diff = diff.view(diff.shape[0], -1) 30 | diff = torch.sum(diff, dim=1) 31 | diff = torch.mean(diff, dim=0) 32 | return diff 33 | 34 | 35 | def L2_each_instance(x: Tensor, y: Tensor) -> Tensor: 36 | """ 37 | N, ?, ?, ?, ... 38 | :return: N, 39 | """ 40 | x = (x - y) ** 2 41 | x = x.view(x.shape[0], -1) 42 | x = torch.norm(x, dim=1, p=2) 43 | return x 44 | 45 | 46 | def list_mean(x: List) -> float: 47 | return sum(x) / len(x) 48 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/PurificationDefense.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | from typing import Callable 4 | 5 | 6 | class PurificationDefense(nn.Module): 7 | def __init__(self, purifier: Callable, classifier: Callable, device=torch.device("cuda")): 8 | super(PurificationDefense, self).__init__() 9 | self.purifier = purifier 10 | self.classifier = classifier 11 | self.device = device 12 | self.eval().requires_grad_(False).to(self.device) 13 | 14 | def forward(self, x: Tensor) -> Tensor: 15 | x = self.purifier(x) 16 | x = self.classifier(x) 17 | return x 18 | -------------------------------------------------------------------------------- /defenses/PurificationDefenses/__init__.py: -------------------------------------------------------------------------------- 1 | from .DiffPure import * 2 | from .PurificationDefense import PurificationDefense 3 | -------------------------------------------------------------------------------- /defenses/RandomizedSmoothing/RSModels.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .core import Smooth 3 | from models import resnet50 4 | from torchvision import transforms 5 | 6 | 7 | class RandomizedSmoothing(torch.nn.Module): 8 | def __init__(self, sigma="0.50", 9 | pretrained=True, 10 | ): 11 | super(RandomizedSmoothing, self).__init__() 12 | self.device = torch.device('cuda') 13 | self.transform = transforms.Compose([ 14 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), 15 | ]) 16 | model = resnet50() 17 | if pretrained: 18 | r = {} 19 | state = torch.load( 20 | f'./resources/checkpoints/RS/models/' 21 | f'imagenet/resnet50/noise_{sigma}/checkpoint.pth.tar')['state_dict'] 22 | for x in list(state.items()): 23 | if len(x) == 2: 24 | k, v = x 25 | r[k[9:]] = v 26 | model.load_state_dict(r) 27 | model.eval().to(self.device) 28 | self.model = Smooth(model, 1000, float(sigma)) 29 | 30 | def __call__(self, xs): 31 | xs = self.transform(xs) 32 | results = [] 33 | for x in xs.split(1, dim=0): 34 | x = x.squeeze() 35 | x = self.model.predict(x, n=100, alpha=0.001, batch_size=10) 36 | result = torch.zeros((1, 1000), device=self.device) 37 | result[0, x] = 1 38 | results.append(result) 39 | results = torch.cat(results, dim=0) 40 | return results 41 | 42 | 43 | def randomized_smoothing_resnet50(sigma="0.00", pretrained=True): 44 | return RandomizedSmoothing(sigma, pretrained=pretrained) 45 | -------------------------------------------------------------------------------- /defenses/RandomizedSmoothing/__init__.py: -------------------------------------------------------------------------------- 1 | from .core import Smooth 2 | from .RSModels import randomized_smoothing_resnet50 3 | -------------------------------------------------------------------------------- /defenses/Transformations/BitDepthReduction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import transforms 3 | 4 | __all__ = ['BitDepthReduction'] 5 | 6 | class BitDepthReductionFunction(object): 7 | ''' 8 | https://github.com/thu-ml/ares/blob/main/pytorch_ares/pytorch_ares/defense_torch/bit_depth_reduction.py 9 | ''' 10 | 11 | def __init__(self, compressed_bit=4): 12 | self.compressed_bit = compressed_bit 13 | self.device = torch.device('cuda') 14 | 15 | def bit_depth_reduction(self, xs): 16 | bits = 2 ** self.compressed_bit # 2**i 17 | xs_compress = (xs.detach() * bits).int() 18 | xs_255 = (xs_compress * (255 / bits)) 19 | xs_compress = (xs_255 / 255).to(self.device) 20 | return xs_compress 21 | 22 | def __call__(self, *args, **kwargs): 23 | return self.bit_depth_reduction(*args, **kwargs) 24 | 25 | 26 | class BitDepthReduction(torch.nn.Module): 27 | def __init__(self, model: torch.nn.Module,): 28 | super(BitDepthReduction, self).__init__() 29 | self.model = model 30 | self.bit = BitDepthReductionFunction() 31 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 32 | 33 | def forward(self, x): 34 | x = self.bit(x) 35 | return self.model(x) 36 | -------------------------------------------------------------------------------- /defenses/Transformations/JPEGCompression.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import transforms 3 | from PIL import Image 4 | from io import BytesIO # BytesIO write bytes in memory 5 | 6 | _to_pil_image = transforms.ToPILImage() 7 | _to_tensor = transforms.ToTensor() 8 | 9 | __all__ = ['JPEGCompression'] 10 | 11 | 12 | class Jpeg_compression(object): 13 | ''' 14 | https://github.com/thu-ml/ares/blob/main/pytorch_ares/pytorch_ares/defense_torch/jpeg_compression.py 15 | ''' 16 | 17 | def __init__(self, quality=75): 18 | self.quality = quality 19 | 20 | def jpegcompression(self, x): 21 | lst_img = [] 22 | for img in x: 23 | img = _to_pil_image(img.detach().clone().cpu()) 24 | virtualpath = BytesIO() 25 | img.save(virtualpath, 'JPEG', quality=self.quality) 26 | lst_img.append(_to_tensor(Image.open(virtualpath))) 27 | return x.new_tensor(torch.stack(lst_img)) 28 | 29 | def __call__(self, *args, **kwargs): 30 | return self.jpegcompression(*args, **kwargs) 31 | 32 | 33 | class JPEGCompression(torch.nn.Module): 34 | def __init__(self, model: torch.nn.Module, 35 | transform=transforms.Compose([ 36 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 37 | ])): 38 | super(JPEGCompression, self).__init__() 39 | self.model = model 40 | self.transforms = transform 41 | self.jpegcompression = Jpeg_compression() 42 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 43 | 44 | def forward(self, x): 45 | x = self.jpegcompression(x) 46 | x = x.to(self.device) 47 | return self.model(x) 48 | -------------------------------------------------------------------------------- /defenses/Transformations/__init__.py: -------------------------------------------------------------------------------- 1 | from .BitDepthReduction import BitDepthReduction 2 | from .JPEGCompression import JPEGCompression 3 | from .Randomization import Randomization 4 | from .TTE import TTEAugWrapper 5 | -------------------------------------------------------------------------------- /defenses/__init__.py: -------------------------------------------------------------------------------- 1 | from .Transformations import BitDepthReduction, Randomization, JPEGCompression 2 | from .NeuralRepresentationPurifier import NeuralRepresentationPurifier 3 | from .RandomizedSmoothing import randomized_smoothing_resnet50 4 | from defenses.PurificationDefenses.DiffPure import * 5 | -------------------------------------------------------------------------------- /experiments/CIFAR10C.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test, get_cifar_10_c_loader 3 | from tester import test_apgd_dlr_acc, test_transfer_attack_acc, test_acc 4 | from torchvision import transforms 5 | import numpy as np 6 | import random 7 | from defenses.PurificationDefenses.DiffPure import RobustDiffusionClassifier 8 | 9 | torch.manual_seed(1) 10 | random.seed(1) 11 | np.random.seed(1) 12 | if torch.cuda.is_available(): 13 | torch.cuda.manual_seed_all(1) 14 | 15 | to_img = transforms.ToPILImage() 16 | corruptions = [ 17 | 'glass_blur', 18 | 'gaussian_noise', 19 | 'shot_noise', 20 | 'speckle_noise', 21 | 'impulse_noise', 22 | 'defocus_blur', 23 | 'gaussian_blur', 24 | 'motion_blur', 25 | 'zoom_blur', 26 | 'snow', 27 | 'fog', 28 | 'brightness', 29 | 'contrast', 30 | 'elastic_transform', 31 | 'pixelate', 32 | 'jpeg_compression', 33 | 'spatter', 34 | 'saturate', 35 | 'frost', 36 | ] 37 | 38 | device = torch.device('cuda') 39 | 40 | 41 | def save_img(x, name='test'): 42 | if x.ndim == 4: 43 | x = x[0] 44 | img = to_img(x) 45 | img.save(f'{name}.png') 46 | 47 | 48 | diffpure = RobustDiffusionClassifier( 49 | False, False, False 50 | ).cuda().requires_grad_(False).eval() 51 | accs = [] 52 | for name in corruptions: 53 | loader = get_cifar_10_c_loader(name=name, batch_size=1) 54 | print(f'now testing **{name}**, total images {len(loader)}') 55 | _, now_acc = test_acc(diffpure, loader) 56 | accs.append(now_acc) 57 | print('-' * 100) 58 | acc = sum(accs) / len(accs) 59 | -------------------------------------------------------------------------------- /experiments/DiffAttack.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks import BIM 3 | from data import get_CIFAR10_test 4 | from tester import test_robustness, test_apgd_dlr_acc 5 | from defenses import DiffusionPure 6 | from utils import set_seed 7 | import argparse 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument("--begin", type=int) 11 | parser.add_argument("--end", type=int) 12 | args = parser.parse_args() 13 | begin, end = args.begin, args.end 14 | 15 | set_seed(1) 16 | loader = get_CIFAR10_test(batch_size=1) 17 | device = torch.device("cuda") 18 | diffpure = DiffusionPure(mode="sde", grad_checkpoint=True).eval().requires_grad_(False) 19 | 20 | loader = [(x, y) for i, ((x, y)) in enumerate(loader) if begin <= i < end] 21 | attacker = BIM([diffpure], step_size=0.1, total_step=200, eot_step=32, epsilon=0.5, norm="L2", eot_batch_size=32) 22 | test_robustness( 23 | attacker, 24 | loader, 25 | [diffpure], 26 | ) 27 | # test_apgd_dlr_acc(diffpure, loader=loader, bs=16, eot_iter=20, norm="L2", eps=0.5) 28 | -------------------------------------------------------------------------------- /experiments/DiffusionAsClassifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test 3 | from tester import test_apgd_dlr_acc 4 | from torchvision import transforms 5 | import numpy as np 6 | import random 7 | from defenses.PurificationDefenses.DiffPure import RobustDiffusionClassifier 8 | import argparse 9 | 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--begin', type=int) 12 | parser.add_argument('--end', type=int) 13 | args = parser.parse_args() 14 | begin, end = args.begin, args.end 15 | print(args.begin, args.end) 16 | 17 | torch.manual_seed(1) 18 | random.seed(1) 19 | np.random.seed(1) 20 | if torch.cuda.is_available(): 21 | torch.cuda.manual_seed_all(1) 22 | 23 | to_img = transforms.ToPILImage() 24 | loader = get_CIFAR10_test(batch_size=1) 25 | device = torch.device('cuda') 26 | 27 | diffpure = RobustDiffusionClassifier( 28 | bpda=True, 29 | likelihood_maximization=True, 30 | diffpure=False, 31 | second_order=False 32 | ) 33 | 34 | diffpure.eval().requires_grad_(False).to(device) 35 | xs, ys = [], [] 36 | for x, y in loader: 37 | xs.append(x) 38 | ys.append(y) 39 | x = torch.concat(xs, dim=0).cuda() 40 | y = torch.concat(ys, dim=0).cuda() 41 | x, y = x[begin:end], y[begin:end] 42 | 43 | test_apgd_dlr_acc(diffpure, x=x, y=y, bs=1, log_path=f'./Linfbpda-{begin}-{end}.txt', 44 | eps=8 / 255, norm='Linf') 45 | -------------------------------------------------------------------------------- /experiments/DiffusionAsClassifierN+T.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_NCSNPP_cached, DiffusionClassifierCached, get_edm_cifar_uncond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 7 | diffusion_likelihood_maximizer_defense, 8 | ) 9 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC 10 | from defenses.PurificationDefenses.DiffPure.LikelihoodMaximization import EDMEulerIntegralLM 11 | from attacks import StAdvAttack 12 | from tester import test_robustness 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument("--begin", type=int) 16 | parser.add_argument("--end", type=int) 17 | args = parser.parse_args() 18 | begin, end = args.begin, args.end 19 | 20 | model = get_NCSNPP_cached(grad_checkpoint=True).cuda() 21 | test_loader = get_CIFAR10_test(batch_size=1) 22 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 23 | model.load_state_dict( 24 | torch.load("/workspace/home/chenhuanran2022/work/Diffusion/" "checkpoints/cached_kd_edm_all/ema_1800.pt") 25 | ) 26 | 27 | edm_unet = get_edm_cifar_uncond() 28 | edm_unet.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 29 | purify_dc = EDMEulerIntegralDC(unet=edm_unet) 30 | lm = EDMEulerIntegralLM(purify_dc) 31 | 32 | classifier = DiffusionClassifierCached(model) 33 | 34 | defensed = diffusion_likelihood_maximizer_defense(classifier, lm.likelihood_maximization_T1) 35 | 36 | # test_apgd_dlr_acc(defensed, loader=test_loader, norm="L2", eps=0.5) 37 | attacker = StAdvAttack(defensed, num_iterations=100, bound=0.05) 38 | test_robustness(attacker, test_loader, [defensed]) 39 | -------------------------------------------------------------------------------- /experiments/DiffusionAsClassifierN+TK.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 7 | diffusion_likelihood_maximizer_defense, 8 | ) 9 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 10 | from defenses.PurificationDefenses.DiffPure.LikelihoodMaximization import EDMEulerIntegralLM 11 | 12 | parser = argparse.ArgumentParser() 13 | parser.add_argument("--begin", type=int) 14 | parser.add_argument("--end", type=int) 15 | args = parser.parse_args() 16 | begin, end = args.begin, args.end 17 | 18 | edm_unet = get_edm_cifar_uncond() 19 | edm_unet.load_state_dict(torch.load("../../resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 20 | purify_dc = EDMEulerIntegralDC(unet=edm_unet) 21 | lm = EDMEulerIntegralLM(purify_dc) 22 | 23 | model = get_edm_cifar_cond().cuda() 24 | test_loader = get_CIFAR10_test(batch_size=1) 25 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 26 | model.load_state_dict(torch.load("../../resources/checkpoints/EDM/edm_cifar_cond.pt")) 27 | 28 | dc = EDMEulerIntegralWraped(unet=model) 29 | defensed = diffusion_likelihood_maximizer_defense(dc, lm.likelihood_maximization_T1) 30 | 31 | test_apgd_dlr_acc(defensed, loader=test_loader, norm="L2", eps=0.5) 32 | -------------------------------------------------------------------------------- /experiments/DiffusionAsClassifierT.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_NCSNPP_cached, DiffusionClassifierCached, get_edm_cifar_uncond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC 7 | from attacks import StAdvAttack 8 | from tester import test_robustness 9 | 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument("--begin", type=int) 12 | parser.add_argument("--end", type=int) 13 | args = parser.parse_args() 14 | begin, end = args.begin, args.end 15 | 16 | 17 | model = get_NCSNPP_cached(grad_checkpoint=True).cuda() 18 | test_loader = get_CIFAR10_test(batch_size=1) 19 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 20 | # actually is ema_3050 21 | model.load_state_dict( 22 | torch.load("/workspace/home/chenhuanran2022/work/Diffusion/" "checkpoints/cached_kd_edm_all/ema_new.pt") 23 | ) 24 | 25 | edm_unet = get_edm_cifar_uncond() 26 | edm_unet.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 27 | purify_dc = EDMEulerIntegralDC(unet=edm_unet) 28 | 29 | classifier = DiffusionClassifierCached(model) 30 | 31 | test_apgd_dlr_acc(classifier, loader=test_loader, norm="Linf", eps=8/255) 32 | test_apgd_dlr_acc(classifier, loader=test_loader, norm="L2", eps=0.5) 33 | # attacker = StAdvAttack(classifier, num_iterations=100, bound=0.05) 34 | # test_robustness(attacker, test_loader, [classifier]) 35 | -------------------------------------------------------------------------------- /experiments/DiffusionAsClassifierTK.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 7 | 8 | parser = argparse.ArgumentParser() 9 | parser.add_argument("--begin", type=int) 10 | parser.add_argument("--end", type=int) 11 | args = parser.parse_args() 12 | begin, end = args.begin, args.end 13 | 14 | model = get_edm_cifar_cond().cuda() 15 | test_loader = get_CIFAR10_test(batch_size=1) 16 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 17 | model.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_cond.pt")) 18 | 19 | dc = EDMEulerIntegralWraped(unet=model) 20 | 21 | test_apgd_dlr_acc(dc, loader=test_loader, norm="L2", eps=0.5) 22 | -------------------------------------------------------------------------------- /experiments/DiffusionMaximizer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test 3 | from tester import test_apgd_dlr_acc 4 | from torchvision import transforms 5 | import numpy as np 6 | import random 7 | from defenses.PurificationDefenses.DiffPure import diffusion_likelihood_maximizer_defense 8 | from models import WideResNet_70_16_dropout 9 | 10 | torch.manual_seed(1) 11 | random.seed(1) 12 | np.random.seed(1) 13 | if torch.cuda.is_available(): 14 | torch.cuda.manual_seed_all(1) 15 | 16 | to_img = transforms.ToPILImage() 17 | loader = get_CIFAR10_test(batch_size=1) 18 | device = torch.device('cuda') 19 | 20 | model = WideResNet_70_16_dropout().eval().requires_grad_(False) 21 | diffpure = diffusion_likelihood_maximizer_defense(model) 22 | 23 | xs, ys = [], [] 24 | for x, y in loader: 25 | xs.append(x) 26 | ys.append(y) 27 | x = torch.concat(xs, dim=0).cuda() 28 | y = torch.concat(ys, dim=0).cuda() 29 | x, y = x[0:100], y[0:100] 30 | test_apgd_dlr_acc(diffpure, x=x, y=y, bs=1, log_path='./maximizer.txt') 31 | -------------------------------------------------------------------------------- /experiments/DirectAttack.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test 3 | from tester import test_apgd_dlr_acc 4 | from torchvision import transforms 5 | import numpy as np 6 | import random 7 | from defenses.PurificationDefenses.DiffPure import RobustDiffusionClassifier 8 | 9 | torch.manual_seed(1) 10 | random.seed(1) 11 | np.random.seed(1) 12 | if torch.cuda.is_available(): 13 | torch.cuda.manual_seed_all(1) 14 | 15 | to_img = transforms.ToPILImage() 16 | loader = get_CIFAR10_test(batch_size=1) 17 | device = torch.device('cuda') 18 | 19 | 20 | def save_img(x, name='test'): 21 | if x.ndim == 4: 22 | x = x[0] 23 | img = to_img(x) 24 | img.save(f'{name}.png') 25 | 26 | 27 | diffpure = RobustDiffusionClassifier(second_order=True, likelihood_maximization=True) 28 | unet = diffpure.unet 29 | diffpure.eval().requires_grad_(False).to(device) 30 | xs, ys = [], [] 31 | for x, y in loader: 32 | xs.append(x) 33 | ys.append(y) 34 | x = torch.concat(xs, dim=0).cuda() 35 | y = torch.concat(ys, dim=0).cuda() 36 | x, y = x[:512], y[:512] 37 | test_apgd_dlr_acc(diffpure, x=x, y=y, bs=1, log_path='./direct.txt') 38 | -------------------------------------------------------------------------------- /experiments/LikelihoodMaximization.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from data import get_CIFAR10_test, get_someset_loader 4 | from tester import test_acc, test_apgd_dlr_acc, test_robustness 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 7 | diffusion_likelihood_maximizer_defense, 8 | ) 9 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC 10 | from defenses.PurificationDefenses.DiffPure.LikelihoodMaximization import EDMEulerIntegralLM 11 | from models import WideResNet_70_16_dropout, LogitEnsemble 12 | from attacks import StAdvAttack, BIM 13 | 14 | loader = get_CIFAR10_test(batch_size=16) 15 | 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument("--begin", type=int) 18 | parser.add_argument("--end", type=int) 19 | args = parser.parse_args() 20 | begin, end = args.begin, args.end 21 | loader = [item for i, item in enumerate(loader) if begin <= i < end] 22 | # loader = get_someset_loader(img_path="./debug/LMnofixnoise", gt_path="./debug/LMfixnoise/labels.npy", batch_size=16) 23 | 24 | uncond_edm_net = get_edm_cifar_uncond() 25 | uncond_edm_net.load_state_dict(torch.load("../../resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 26 | uncond_edm_dc = EDMEulerIntegralDC(unet=uncond_edm_net) 27 | lm = EDMEulerIntegralLM(uncond_edm_dc) 28 | 29 | 30 | defensed = diffusion_likelihood_maximizer_defense( 31 | WideResNet_70_16_dropout(), 32 | lm.likelihood_maximization_T1, 33 | ) 34 | 35 | 36 | defensed.eval().requires_grad_(False) 37 | # result = [] 38 | # for i in tqdm(range(10)): 39 | # result.extend(lm.SDS_T1(loader[i][0].cuda(), return_intermediate_result=True, eps=300, iter_step=20)[1]) 40 | # concatenate_image(result, col=20, row=10) 41 | test_acc(defensed, loader) 42 | test_apgd_dlr_acc(defensed, loader=loader, norm="Linf", eps=8 / 255, bs=64) 43 | test_apgd_dlr_acc(defensed, loader=loader, norm="L2", eps=0.5, bs=64) 44 | attacker = StAdvAttack(defensed, num_iterations=100, bound=0.05) 45 | test_robustness(attacker, loader, [defensed]) 46 | -------------------------------------------------------------------------------- /experiments/ObfuscateGradient/ObfuscateGradient.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test 3 | from tqdm import tqdm 4 | from torchvision import transforms 5 | import numpy as np 6 | import random 7 | from defenses import RobustDiffusionClassifier, DiffusionClassifier, DiffusionPure 8 | from tester.utils import cosine_similarity 9 | 10 | torch.manual_seed(1) 11 | random.seed(1) 12 | np.random.seed(1) 13 | if torch.cuda.is_available(): 14 | torch.cuda.manual_seed_all(1) 15 | 16 | to_img = transforms.ToPILImage() 17 | loader = get_CIFAR10_test(batch_size=1) 18 | device = torch.device('cuda') 19 | 20 | diffpure = DiffusionClassifier() 21 | diffpure.unet.load_state_dict(torch.load('./unet_condition_old.pt')) 22 | diffpure.eval().requires_grad_(False).to(device) 23 | 24 | x, y = next(iter(loader)) 25 | x, y = x.cuda(), y.cuda() 26 | criterion = torch.nn.CrossEntropyLoss() 27 | 28 | num_grad = 10 29 | 30 | # Our diffusion classifier 31 | grad = [] 32 | for i in range(10): 33 | grad.clear() 34 | for _ in tqdm(range(num_grad)): 35 | x.requires_grad_(True) 36 | diffpure.partial(x, i) 37 | grad.append(x.grad.clone()) 38 | x.requires_grad_(False) 39 | x.grad = None 40 | 41 | print('class ', i, ' cosine similarity: ', cosine_similarity(grad)) 42 | 43 | # robust diffusion classifier 44 | diffpure = RobustDiffusionClassifier(bpda=True, likelihood_maximization=True, diffpure=False) 45 | diffpure.unet.load_state_dict(torch.load('./unet_condition_old.pt')) 46 | diffpure.eval().requires_grad_(False).cuda() 47 | grad = [] 48 | for i in range(10): 49 | grad.clear() 50 | for _ in tqdm(range(num_grad)): 51 | x.requires_grad_(True) 52 | loss = diffpure(x).squeeze()[i] 53 | loss.backward() 54 | grad.append(x.grad.clone()) 55 | x.requires_grad_(False) 56 | x.grad = None 57 | 58 | print('class ', i, ' cosine similarity: ', cosine_similarity(grad)) 59 | 60 | # diffpure 61 | grad.clear() 62 | del diffpure 63 | criterion = torch.nn.CrossEntropyLoss() 64 | diffpure = DiffusionPure(grad_checkpoint=True) 65 | for _ in tqdm(range(num_grad)): 66 | x.requires_grad_(True) 67 | pre = diffpure(x) 68 | loss = pre[0, 0] 69 | loss.backward() 70 | grad.append(x.grad.clone()) 71 | x.requires_grad_(False) 72 | x.grad = None 73 | 74 | print(cosine_similarity(grad)) 75 | 76 | # resnet 77 | from models import resnet50 78 | 79 | grad.clear() 80 | diffpure = resnet50(pretrained=True).cuda().eval().requires_grad_(False) 81 | x = x.view(1, 3, 32, 32) 82 | for _ in tqdm(range(num_grad)): 83 | x.requires_grad_(True) 84 | pre = diffpure(x) 85 | loss = pre[0, 0] 86 | loss.backward() 87 | grad.append(x.grad.clone()) 88 | x.requires_grad_(False) 89 | x.grad = None 90 | 91 | print(cosine_similarity(grad)) 92 | -------------------------------------------------------------------------------- /experiments/ObfuscateGradient/ObfuscateGradient_likelihoodmaximizer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | from tester.utils import cosine_similarity 6 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 7 | diffusion_likelihood_maximizer_defense, 8 | ) 9 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC 10 | from defenses.PurificationDefenses.DiffPure.LikelihoodMaximization import EDMEulerIntegralLM 11 | from models import WideResNet_70_16_dropout 12 | from tqdm import tqdm 13 | 14 | 15 | loader = get_CIFAR10_test(batch_size=1) 16 | x, y = next(iter(loader)) 17 | x, y = x.cuda(), y.cuda() 18 | uncond_edm_net = get_edm_cifar_uncond() 19 | uncond_edm_net.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 20 | uncond_edm_dc = EDMEulerIntegralDC(unet=uncond_edm_net) 21 | lm = EDMEulerIntegralLM(uncond_edm_dc) 22 | 23 | # cond_edm_net = get_edm_cifar_cond() 24 | # cond_edm_net.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_cond.pt")) 25 | # dc = EDMEulerIntegralDC(unet=cond_edm_net) 26 | 27 | defensed = diffusion_likelihood_maximizer_defense(WideResNet_70_16_dropout(), lm.optimize_back) 28 | 29 | num_grad = 10 30 | 31 | # Our diffusion classifier 32 | criterion = torch.nn.CrossEntropyLoss() 33 | grad = [] 34 | for i in range(10): 35 | grad.clear() 36 | for _ in tqdm(range(num_grad)): 37 | x.requires_grad_(True) 38 | loss = defensed(x).squeeze()[i] 39 | loss.backward() 40 | grad.append(x.grad.clone()) 41 | x.requires_grad_(False) 42 | x.grad = None 43 | 44 | print("class ", i, " cosine similarity: ", cosine_similarity(grad)) 45 | -------------------------------------------------------------------------------- /experiments/ObfuscateGradient/Obfuscate_N+TK.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from models.unets import get_edm_cifar_uncond, get_edm_cifar_cond 3 | from data import get_CIFAR10_test 4 | from tester import test_acc, test_apgd_dlr_acc 5 | import argparse 6 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 7 | diffusion_likelihood_maximizer_defense, 8 | ) 9 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 10 | 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument("--begin", type=int) 13 | parser.add_argument("--end", type=int) 14 | args = parser.parse_args() 15 | begin, end = args.begin, args.end 16 | 17 | edm_unet = get_edm_cifar_uncond() 18 | edm_unet.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 19 | purify_dc = EDMEulerIntegralDC(unet=edm_unet) 20 | 21 | model = get_edm_cifar_cond().cuda() 22 | test_loader = get_CIFAR10_test(batch_size=1) 23 | test_loader = [item for i, item in enumerate(test_loader) if begin <= i < end] 24 | model.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_cond.pt")) 25 | 26 | from tqdm import tqdm 27 | from tester.utils import cosine_similarity 28 | from models import WideResNet_70_16_dropout 29 | from attacks import BIM 30 | 31 | dc = EDMEulerIntegralWraped(unet=model) 32 | defensed = diffusion_likelihood_maximizer_defense(WideResNet_70_16_dropout(), purify_dc) 33 | num_grad = 10 34 | 35 | x, y = test_loader[0] 36 | x, y = x.cuda(), y.cuda() 37 | attacker = BIM([defensed], step_size=1 / 255, total_step=1000) 38 | x = attacker(x, y) 39 | grad = [] 40 | for i in range(10): 41 | grad.clear() 42 | for _ in tqdm(range(num_grad)): 43 | x.requires_grad_(True) 44 | loss = defensed(x).squeeze()[i] 45 | loss.backward() 46 | grad.append(x.grad.clone()) 47 | x.requires_grad_(False) 48 | x.grad = None 49 | 50 | print("class ", i, " cosine similarity: ", cosine_similarity(grad)) 51 | # test_apgd_dlr_acc(dc, loader=test_loader, norm='L2', eps=0.5) 52 | -------------------------------------------------------------------------------- /experiments/OptimaDiffusionClassifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test 3 | from tester import test_apgd_dlr_acc 4 | from torchvision import transforms 5 | import numpy as np 6 | import random 7 | from defenses import OptimalDiffusionClassifier 8 | 9 | torch.manual_seed(1) 10 | random.seed(1) 11 | np.random.seed(1) 12 | if torch.cuda.is_available(): 13 | torch.cuda.manual_seed_all(1) 14 | 15 | to_img = transforms.ToPILImage() 16 | loader = get_CIFAR10_test(batch_size=1) 17 | device = torch.device('cuda') 18 | 19 | loader = [(x, y) for i, ((x, y)) in enumerate(loader) if i < 100] 20 | diffpure = OptimalDiffusionClassifier(loader) 21 | test_apgd_dlr_acc(diffpure, loader=loader) 22 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/ATImageNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_someset_loader 3 | from tester import test_acc, test_apgd_dlr_acc 4 | from utils.seed import set_seed 5 | from models import Wong2020Fast, ClassSelectionModel, RegionClassSelectionModel 6 | from models.RobustBench.imagenet import * 7 | from torch.nn import functional as F 8 | 9 | set_seed(1) 10 | loader = get_someset_loader( 11 | "./resources/RestrictedImageNet256", 12 | "./resources/RestrictedImageNet256/gt.npy", 13 | batch_size=32, 14 | ) 15 | # loader = get_restricted_imagenet_test_loader(batch_size=256, shuffle=True) 16 | x, y = next(iter(loader)) 17 | x, y = x.cuda(), y.cuda() 18 | # save_dataset(x, y, path='./resources/RestrictedImageNet256', gt_saving_name='gt.npy') 19 | models = [Engstrom2019Robustness(), Wong2020Fast(), Salman2020Do_R50(), Debenedetti2022Light_XCiT_S12()] 20 | for model in models: 21 | dc = RegionClassSelectionModel(model, 22 | target_class=RegionClassSelectionModel.restricted_imagenet_target_class() 23 | ).cuda() 24 | # test_acc(dc, loader) 25 | test_apgd_dlr_acc(dc, loader=loader, eps=4 / 255, bs=32) 26 | # for model in models: 27 | # dc = ClassSelectionModel(model, target_class=(151, 281, 30, 33, 80, 365, 389, 118, 300)).cuda() 28 | # x.requires_grad = True 29 | # loss = F.cross_entropy(dc(x), y) 30 | # loss.backward() 31 | # print(torch.mean(x.grad.abs())) 32 | # x.requires_grad = False 33 | # x.grad = None 34 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/DiffAttackImageNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks import BIM 3 | from data import get_someset_loader, get_NIPS17_loader, get_CIFAR10_test 4 | from tester import test_robustness, test_acc 5 | from defenses import DiffusionPure 6 | from utils.seed import set_seed 7 | from models import BaseNormModel, resnet50, RegionClassSelectionModel 8 | from models.unets.guided_diffusion import get_guided_diffusion_unet 9 | from torchvision import transforms 10 | 11 | set_seed(1) 12 | 13 | loader = get_someset_loader( 14 | "./resources/RestrictedImageNet256", 15 | "./resources/RestrictedImageNet256/gt.npy", 16 | batch_size=1, 17 | transform=transforms.Compose( 18 | [ 19 | transforms.Resize((128, 128)), 20 | transforms.ToTensor(), 21 | ] 22 | ), 23 | ) 24 | 25 | device = torch.device("cuda") 26 | unet = get_guided_diffusion_unet(resolution=128) 27 | 28 | diffpure = DiffusionPure( 29 | mode="sde", 30 | grad_checkpoint=True, 31 | unet=unet, 32 | model=RegionClassSelectionModel( 33 | BaseNormModel(resnet50(pretrained=True)), 34 | target_class=RegionClassSelectionModel.restricted_imagenet_target_class(), 35 | ), 36 | img_shape=(3, 128, 128), 37 | post_transforms=transforms.Resize((256, 256)), 38 | ) 39 | diffpure.eval().requires_grad_(False) 40 | 41 | loader = [item for i, item in enumerate(loader) if 0 <= i < 32] 42 | 43 | attacker = BIM([diffpure], step_size=1 / 255, total_step=80, 44 | eot_step=20, epsilon=4 / 255, norm="Linf", eot_batch_size=4) 45 | test_robustness( 46 | attacker, 47 | loader, 48 | [diffpure], 49 | ) 50 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/EDMDCImageNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_someset_loader 3 | from tester import test_acc, test_apgd_dlr_acc 4 | from models.unets.EDM import get_edm_imagenet_64x64_cond 5 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 6 | from utils.seed import set_seed 7 | from torchvision import transforms 8 | from defenses.PurificationDefenses.DiffPure.LikelihoodMaximization import EDMEulerIntegralLM 9 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 10 | diffusion_likelihood_maximizer_defense, 11 | ) 12 | import argparse 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument("--begin", type=int) 16 | parser.add_argument("--end", type=int) 17 | args = parser.parse_args() 18 | begin, end = args.begin, args.end 19 | 20 | set_seed(1) 21 | loader = get_someset_loader( 22 | "./resources/RestrictedImageNet256", 23 | "./resources/RestrictedImageNet256/gt.npy", 24 | batch_size=1, 25 | shuffle=False, 26 | ) 27 | # loader = get_restricted_imagenet_test_loader(batch_size=256, shuffle=True) 28 | # save_dataset(x, y, path='./resources/RestrictedImageNet256', gt_saving_name='gt.npy') 29 | loader = [item for i, item in enumerate(loader) if begin <= i < end] 30 | unet = get_edm_imagenet_64x64_cond() 31 | unet.load_state_dict(torch.load("../../../resources/checkpoints/EDM/edm-imagenet-64x64-cond.pt")) 32 | resizer = transforms.Resize((64, 64), antialias=True) 33 | 34 | 35 | def edm_64x64_transform(x): 36 | x = resizer(x) 37 | x = (x - 0.5) * 2 38 | return x 39 | 40 | dc = EDMEulerIntegralWraped( 41 | unet, target_class=(151, 281, 30, 33, 80, 365, 389, 118, 300), num_classes=1000, transform=edm_64x64_transform 42 | ) 43 | test_apgd_dlr_acc(dc, loader=loader, eps=4 / 255) 44 | # dc = EDMEulerIntegralWraped( 45 | # unet=unet, 46 | # target_class=(151, 281, 30, 33, 80, 365, 389, 118, 300), 47 | # num_classes=1000, 48 | # transform=edm_64x64_transform, 49 | # timesteps=torch.tensor([1.0, 2.0]), 50 | # ) 51 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/EDMLMImageNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_someset_loader 3 | from tester import test_acc, test_apgd_dlr_acc 4 | from models.unets import get_guided_diffusion_unet 5 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 6 | from utils.seed import set_seed 7 | from torchvision import transforms 8 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralLM, VP2EDM 9 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 10 | diffusion_likelihood_maximizer_defense, 11 | ) 12 | import argparse 13 | from models import Salman2020Do_R50, RegionClassSelectionModel 14 | 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument("--begin", type=int) 17 | parser.add_argument("--end", type=int) 18 | args = parser.parse_args() 19 | begin, end = args.begin, args.end 20 | 21 | set_seed(1) 22 | loader = get_someset_loader( 23 | "./resources/RestrictedImageNet256", 24 | "./resources/RestrictedImageNet256/gt.npy", 25 | batch_size=1, 26 | shuffle=False, 27 | transform=transforms.Compose( 28 | [ 29 | transforms.Resize((256, 256)), 30 | transforms.ToTensor(), 31 | ] 32 | ), 33 | ) 34 | loader = [item for i, item in enumerate(loader) if begin <= i < end] 35 | unet = get_guided_diffusion_unet(resolution=256) 36 | 37 | dc = EDMEulerIntegralDC(VP2EDM(unet)) 38 | lm = EDMEulerIntegralLM(dc) 39 | classifier = RegionClassSelectionModel( 40 | Salman2020Do_R50(pretrained=True), 41 | target_class=RegionClassSelectionModel.restricted_imagenet_target_class(), 42 | ) 43 | 44 | defensed = diffusion_likelihood_maximizer_defense( 45 | classifier, lambda x: lm.likelihood_maximization_T1(x, iter_step=5, t_min=0.4, t_max=0.6) 46 | ) 47 | # test_acc(dc, loader) 48 | test_apgd_dlr_acc(defensed, loader=loader, eps=4 / 255, bs=2) 49 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/EDMRDCImageNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_someset_loader 3 | from tester import test_acc, test_apgd_dlr_acc 4 | from models.unets import get_guided_diffusion_unet, get_edm_imagenet_64x64_cond 5 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralDC, EDMEulerIntegralWraped 6 | from utils.seed import set_seed 7 | from torchvision import transforms 8 | from defenses.PurificationDefenses.DiffPure import EDMEulerIntegralLM, VP2EDM 9 | from defenses.PurificationDefenses.DiffPure.DiffPure.DiffusionLikelihoodMaximizer import ( 10 | diffusion_likelihood_maximizer_defense, 11 | ) 12 | import argparse 13 | from models import BaseNormModel, resnet50, RegionClassSelectionModel 14 | 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument("--begin", type=int) 17 | parser.add_argument("--end", type=int) 18 | args = parser.parse_args() 19 | begin, end = args.begin, args.end 20 | 21 | set_seed(1) 22 | loader = get_someset_loader( 23 | "./resources/RestrictedImageNet256", 24 | "./resources/RestrictedImageNet256/gt.npy", 25 | batch_size=1, 26 | shuffle=False, 27 | transform=transforms.Compose( 28 | [ 29 | transforms.Resize((256, 256)), 30 | transforms.ToTensor(), 31 | ] 32 | ), 33 | ) 34 | loader = [item for i, item in enumerate(loader) if begin <= i < end] 35 | unet = get_guided_diffusion_unet(resolution=256) 36 | 37 | dc = EDMEulerIntegralDC(VP2EDM(unet)) 38 | lm = EDMEulerIntegralLM(dc) 39 | 40 | edm_unet = get_edm_imagenet_64x64_cond() 41 | edm_unet.load_state_dict(torch.load("../../../resources/checkpoints/EDM/edm-imagenet-64x64-cond.pt")) 42 | 43 | 44 | def edm_64x64_transform(x): 45 | x = transforms.Resize((64, 64))(x) 46 | x = (x - 0.5) * 2 47 | return x 48 | 49 | 50 | dc_wraped = EDMEulerIntegralWraped( 51 | edm_unet, 52 | target_class=(151, 281, 30, 33, 80, 365, 389, 118, 300), 53 | num_classes=1000, 54 | transform=edm_64x64_transform 55 | ) 56 | defensed = diffusion_likelihood_maximizer_defense( 57 | dc_wraped, lambda x: lm.likelihood_maximization_T1(x, iter_step=5, t_min=0.4, t_max=0.6) 58 | ) 59 | # test_acc(dc, loader) 60 | test_apgd_dlr_acc(defensed, loader=loader, eps=4 / 255, bs=1) 61 | 62 | # 目前的策略是，虽然edm没有uncond，但我似然仍然可以用uncond，因为他似乎能泛化到uncond 63 | # 还有个策略就是看下embedding，我能不能取这些class embedding的平均值。 64 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/GuidedDiffusionClassifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from data import get_CIFAR10_test 3 | from models.unets import get_NCSNPP 4 | from tester import test_apgd_dlr_acc, test_transfer_attack_acc 5 | from torchvision import transforms 6 | import numpy as np 7 | import random 8 | from defenses.PurificationDefenses.DiffPure import RobustDiffusionClassifier 9 | import argparse 10 | 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument('--begin', type=int) 13 | parser.add_argument('--end', type=int) 14 | args = parser.parse_args() 15 | begin, end = args.begin, args.end 16 | print(args.begin, args.end) 17 | 18 | torch.manual_seed(1) 19 | random.seed(1) 20 | np.random.seed(1) 21 | if torch.cuda.is_available(): 22 | torch.cuda.manual_seed_all(1) 23 | 24 | to_img = transforms.ToPILImage() 25 | loader = get_CIFAR10_test(batch_size=1) 26 | device = torch.device('cuda') 27 | 28 | unet = get_NCSNPP() 29 | diffpure = RobustDiffusionClassifier( 30 | bpda=False, 31 | likelihood_maximization=False, 32 | diffpure=False, 33 | second_order=False, 34 | unet=unet, 35 | ) 36 | 37 | diffpure.eval().requires_grad_(False).to(device) 38 | xs, ys = [], [] 39 | for x, y in loader: 40 | xs.append(x) 41 | ys.append(y) 42 | x = torch.concat(xs, dim=0).cuda() 43 | y = torch.concat(ys, dim=0).cuda() 44 | x, y = x[begin:end], y[begin:end] 45 | loader = [item for i, item in enumerate(loader) if begin <= i < end] 46 | 47 | target = torch.tensor([[0.1200, 0.1195, 0.1197, 0.1195, 0.1199, 0.1199, 0.1198, 0.1211, 0.1198, 0.1198]], device=device) 48 | ce = torch.nn.CrossEntropyLoss() 49 | 50 | 51 | # def loss(x, y): 52 | # return -torch.mean((x - target.repeat(x.shape[0], 1)) ** 2) 53 | def loss(x, y): 54 | return -torch.mean((x - target.repeat(x.shape[0], 1)) ** 2) + ce(x, y) 55 | 56 | 57 | # 58 | # test_apgd_dlr_acc(diffpure, x=x, y=y, bs=1, log_path=f'./Linfbpda-{begin}-{end}.txt', 59 | # eps=8 / 255, norm='Linf') 60 | from attacks import PGD 61 | 62 | attacker = PGD([diffpure], epsilon=8 / 255, total_step=20, step_size=1 / 255, criterion=loss, 63 | eot_step=1, eot_batch_size=1) 64 | test_transfer_attack_acc(attacker, loader, [diffpure], ) 65 | -------------------------------------------------------------------------------- /experiments/RestrictedImageNet/TTEImageNet.py: -------------------------------------------------------------------------------- 1 | from robustbench import load_model 2 | from defenses.Transformations.TTE import TTEAugWrapper 3 | from data import get_CIFAR10_test 4 | from tester import test_apgd_dlr_acc, test_transfer_attack_acc 5 | from attacks import StAdvAttack 6 | 7 | loader = get_CIFAR10_test(batch_size=1) 8 | loader = [item for i, item in enumerate(loader) if i < 256] 9 | models = [load_model(model_name="Wang2023Better_WRN-70-16", dataset="cifar10", threat_model="Linf"), 10 | load_model(model_name="Wang2023Better_WRN-70-16", dataset="cifar10", threat_model="L2"), 11 | ] 12 | for model in models: 13 | wraped_model = TTEAugWrapper(model, flip=True, flip_crop=True, n_crops=1) 14 | test_apgd_dlr_acc(wraped_model, loader=loader, bs=1, eps=8 / 255, norm='Linf') 15 | test_apgd_dlr_acc(wraped_model, loader=loader, bs=1, eps=0.5, norm='L2') 16 | 17 | # attacker = StAdvAttack([wraped_model], num_iterations=100, bound=0.05) 18 | # test_transfer_attack_acc(attacker, loader, [wraped_model]) 19 | -------------------------------------------------------------------------------- /experiments/cifar100/CIFAR100DCNaN.py: -------------------------------------------------------------------------------- 1 | from defenses.PurificationDefenses.DiffPure import ( 2 | EDMEulerIntegralDC, 3 | EDMEulerIntegralLM, 4 | diffusion_likelihood_maximizer_defense, 5 | DiffusionClassifierSingleHeadBaseWraped, 6 | ) 7 | from models.unets import get_edm_cifar_unet 8 | import torch 9 | from data import get_CIFAR100_test 10 | from tester import test_acc, test_apgd_dlr_acc 11 | import argparse 12 | from utils.seed import set_seed 13 | from torch.amp import autocast 14 | 15 | set_seed(1) 16 | 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument("--begin", type=int) 19 | parser.add_argument("--end", type=int) 20 | parser.add_argument("--t_steps", type=int, default=126) 21 | args = parser.parse_args() 22 | begin, end, t_steps = args.begin, args.end, args.t_steps 23 | 24 | loader = get_CIFAR100_test(batch_size=1) 25 | loader = [item for i, item in enumerate(loader) if begin <= i < end] 26 | device = torch.device("cuda") 27 | 28 | torch.autograd.set_detect_anomaly(True) 29 | cond_edm = get_edm_cifar_unet(dataset="cifar100", cond=True, use_fp16=True) 30 | classifier = EDMEulerIntegralDC(cond_edm, timesteps=torch.linspace(1e-4, 3, steps=t_steps), num_classes=100) 31 | classifier.share_noise = True 32 | classifier = DiffusionClassifierSingleHeadBaseWraped(classifier) 33 | classifier.eval().requires_grad_(False).to(device) 34 | 35 | 36 | def nan_hook(module, f_in, f_out): 37 | # print(module) 38 | # print(torch.max(f_in[0]), torch.max(f_out[0])) 39 | func = torch.isnan 40 | if all([(torch.sum(func(i)) <= 0).item() for i in f_in]) and any([(torch.sum(func(i)) > 0).item() for i in f_out]): 41 | print("nan!!!!!!") 42 | print(module, [(torch.sum(func(i)) <= 0).item() for i in f_in], [(torch.sum(func(i)) > 0).item() for i in f_out]) 43 | assert False 44 | func = torch.isinf 45 | if all([(torch.sum(func(i)) <= 0).item() for i in f_in]) and any([(torch.sum(func(i)) > 0).item() for i in f_out]): 46 | print("inf!!!!") 47 | print(module, [(torch.sum(func(i)) <= 0).item() for i in f_in], [(torch.sum(func(i)) > 0).item() for i in f_out]) 48 | assert False 49 | 50 | 51 | for submodule in classifier.modules(): 52 | submodule.register_forward_hook(nan_hook) 53 | 54 | test_acc(classifier, loader, verbose=True) 55 | # g = Smooth(classifier, batch_size=32, verbose=False) 56 | # result, nAs, ns, radii = certify_robustness(g, loader, n0=10, n=1000) 57 | -------------------------------------------------------------------------------- /experiments/cifar100/CIFAR100DiffPure.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from attacks import BIM 3 | from models import BaseNormModel 4 | from data import get_CIFAR100_test 5 | from tester import test_robustness, test_apgd_dlr_acc, test_acc 6 | from models.unets import get_edm_cifar_unet 7 | from models.SmallResolutionModel import wrn_40_2 8 | from defenses import DiffusionPure 9 | from utils import set_seed 10 | import argparse 11 | from torchvision import transforms 12 | 13 | parser = argparse.ArgumentParser() 14 | parser.add_argument("--begin", type=int) 15 | parser.add_argument("--end", type=int) 16 | args = parser.parse_args() 17 | begin, end = args.begin, args.end 18 | 19 | set_seed(1) 20 | loader = get_CIFAR100_test(batch_size=1) 21 | device = torch.device("cuda") 22 | unet = get_edm_cifar_unet(pretrained=True, dataset="cifar100", cond=False) 23 | model = wrn_40_2(num_classes=100) 24 | model.load_state_dict(torch.load("../../../resources/checkpoints/cifar100/wrn_40_2.pth")["model"]) 25 | model = BaseNormModel(model, transform=transforms.Normalize([0.5071, 0.4867, 0.4408], [0.2675, 0.2565, 0.2761])) 26 | diffpure = DiffusionPure(mode="edm", grad_checkpoint=True, unet=unet, model=model).eval().requires_grad_(False) 27 | 28 | loader = [(x, y) for i, ((x, y)) in enumerate(loader) if begin <= i < end] 29 | test_acc(diffpure, loader) 30 | # attacker = BIM([diffpure], step_size=0.1, total_step=200, eot_step=32, epsilon=0.5, norm="L2", eot_batch_size=32) 31 | # test_robustness( 32 | # attacker, 33 | # loader, 34 | # [diffpure], 35 | # ) 36 | # test_apgd_dlr_acc(diffpure, loader=loader, bs=16, eot_iter=20, norm="L2", eps=0.5) 37 | -------------------------------------------------------------------------------- /experiments/stadv.py: -------------------------------------------------------------------------------- 1 | from robustbench import load_model 2 | import torch 3 | from data import get_CIFAR10_test 4 | from tester import test_transfer_attack_acc 5 | import numpy as np 6 | import random 7 | from attacks.stadv_eot.attacks import StAdvAttack 8 | from defenses import RobustDiffusionClassifier 9 | 10 | torch.manual_seed(1) 11 | random.seed(1) 12 | np.random.seed(1) 13 | if torch.cuda.is_available(): 14 | torch.cuda.manual_seed_all(1) 15 | 16 | loader = get_CIFAR10_test(batch_size=1) 17 | 18 | models = [load_model(model_name="Rebuffi2021Fixing_70_16_cutmix_ddpm", dataset="cifar10", threat_model="L2"), 19 | load_model(model_name="Wang2023Better_WRN-70-16", dataset="cifar10", threat_model="L2"), 20 | load_model(model_name="Rebuffi2021Fixing_70_16_cutmix_ddpm", dataset="cifar10", threat_model="Linf"), 21 | load_model(model_name="Wang2023Better_WRN-70-16", dataset="cifar10", threat_model="Linf")] 22 | # 23 | loader = [(x, y) for i, ((x, y)) in enumerate(loader) if 96 <= i < 128] 24 | 25 | result = [] 26 | for model in models: 27 | model.eval().requires_grad_(False).cuda() 28 | attacker = StAdvAttack(model, num_iterations=100, bound=0.05) 29 | result.append(test_transfer_attack_acc(attacker, loader, [model])) 30 | print(result) 31 | 32 | model = RobustDiffusionClassifier(False, False, False).eval().requires_grad_(False).cuda() 33 | attacker = StAdvAttack(model, num_iterations=100, bound=0.05) 34 | test_transfer_attack_acc(attacker, loader, [model]) 35 | -------------------------------------------------------------------------------- /models/RobustBench/__init__.py: -------------------------------------------------------------------------------- 1 | from .imagenet import Wong2020Fast, Engstrom2019Robustness, \ 2 | Salman2020Do_R50, Salman2020Do_R18, Salman2020Do_50_2, \ 3 | Debenedetti2022Light_XCiT_S12, Debenedetti2022Light_XCiT_M12, Debenedetti2022Light_XCiT_L12 4 | from .cifar10 import * 5 | -------------------------------------------------------------------------------- /models/RobustBench/cifar10.py: -------------------------------------------------------------------------------- 1 | from robustbench.utils import load_model 2 | from robustbench.model_zoo.cifar10 import cifar_10_models 3 | 4 | 5 | def Andriushchenko2020Understanding(pretrained=True): 6 | return load_model(model_name="Andriushchenko2020Understanding", dataset="cifar10", threat_model="Linf") 7 | 8 | 9 | def Carmon2019Unlabeled(pretrained=True): 10 | return load_model(model_name="Carmon2019Unlabeled", dataset="cifar10", threat_model="Linf") 11 | 12 | 13 | def Sehwag2020Hydra(pretrained=True): 14 | return load_model(model_name="Sehwag2020Hydra", dataset="cifar10", threat_model="Linf") 15 | 16 | 17 | def Wang2020Improving(pretrained=True): 18 | return load_model(model_name="Wang2020Improving", dataset="cifar10", threat_model="Linf") 19 | 20 | 21 | def Hendrycks2019Using(pretrained=True): 22 | return load_model(model_name="Hendrycks2019Using", dataset="cifar10", threat_model="Linf") 23 | 24 | 25 | def Wang2023Better(pretrained=True, dataset="cifar10", mode="Linf"): 26 | return load_model(model_name="Wang2023Better_WRN-70-16", dataset=dataset, threat_model=mode) 27 | 28 | 29 | def Rebuffi2021Fixing(pretrained=True, dataset="cifar10", mode="Linf"): 30 | return load_model(model_name="Rebuffi2021Fixing_70_16_cutmix_ddpm", dataset=dataset, threat_model=mode) 31 | 32 | 33 | ################## Below Are L2 Models ########################################## 34 | 35 | 36 | def Rice2020OverfittingNetL2(pretrained=True): 37 | return load_model(model_name="Rice2020Overfitting", dataset="cifar10", threat_model="L2") 38 | -------------------------------------------------------------------------------- /models/RobustBench/imagenet.py: -------------------------------------------------------------------------------- 1 | from robustbench.utils import load_model 2 | from robustbench.model_zoo.imagenet import normalize_model 3 | 4 | 5 | def Wong2020Fast(pretrained=True): 6 | return load_model(model_name="Wong2020Fast", dataset="imagenet", threat_model="Linf") 7 | 8 | 9 | def Engstrom2019Robustness(pretrained=True): 10 | return load_model(model_name="Engstrom2019Robustness", dataset="imagenet", threat_model="Linf") 11 | 12 | 13 | def Salman2020Do_R50(pretrained=True): 14 | return load_model(model_name="Salman2020Do_R50", dataset="imagenet", threat_model="Linf") 15 | 16 | 17 | def Salman2020Do_R18(pretrained=True): 18 | return load_model(model_name="Salman2020Do_R18", dataset="imagenet", threat_model="Linf") 19 | 20 | 21 | def Salman2020Do_50_2(pretrained=True): 22 | return load_model(model_name="Salman2020Do_50_2", dataset="imagenet", threat_model="Linf") 23 | 24 | 25 | def Debenedetti2022Light_XCiT_S12(pretrained=True): 26 | return load_model(model_name="Debenedetti2022Light_XCiT-S12", dataset="imagenet", threat_model="Linf") 27 | 28 | 29 | def Debenedetti2022Light_XCiT_M12(pretrained=True): 30 | return load_model(model_name="Debenedetti2022Light_XCiT-M12", dataset="imagenet", threat_model="Linf") 31 | 32 | 33 | def Debenedetti2022Light_XCiT_L12(pretrained=True): 34 | return load_model(model_name="Debenedetti2022Light_XCiT-L12", dataset="imagenet", threat_model="Linf") -------------------------------------------------------------------------------- /models/SmallResolutionModel/__init__.py: -------------------------------------------------------------------------------- 1 | from .mobilenetv2 import * 2 | from .resnet import * 3 | from .resnet_imagenet import * 4 | from .resnetv2 import * 5 | from .ShuffleNetv1 import * 6 | from .ShuffleNetv2 import * 7 | from .vgg import * 8 | from .wrn import * 9 | from .cifar10_resnet import WideResNet_70_16, WideResNet_70_16_dropout 10 | from .rs_cifar_models import rs_cifar_resnet20, rs_cifar_resnet110 11 | -------------------------------------------------------------------------------- /models/SmallResolutionModel/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.init as init 6 | 7 | 8 | def cal_param_size(model): 9 | return sum([i.numel() for i in model.parameters()]) 10 | 11 | 12 | count_ops = 0 13 | 14 | 15 | def measure_layer(layer, x, multi_add=1): 16 | delta_ops = 0 17 | type_name = str(layer)[: str(layer).find("(")].strip() 18 | 19 | if type_name in ["Conv2d"]: 20 | out_h = int( 21 | (x.size()[2] + 2 * layer.padding[0] - layer.kernel_size[0]) // layer.stride[0] + 1 22 | ) 23 | out_w = int( 24 | (x.size()[3] + 2 * layer.padding[1] - layer.kernel_size[1]) // layer.stride[1] + 1 25 | ) 26 | delta_ops = ( 27 | layer.in_channels 28 | * layer.out_channels 29 | * layer.kernel_size[0] 30 | * layer.kernel_size[1] 31 | * out_h 32 | * out_w 33 | // layer.groups 34 | * multi_add 35 | ) 36 | 37 | ### ops_linear 38 | elif type_name in ["Linear"]: 39 | weight_ops = layer.weight.numel() * multi_add 40 | bias_ops = 0 41 | delta_ops = weight_ops + bias_ops 42 | 43 | global count_ops 44 | count_ops += delta_ops 45 | return 46 | 47 | 48 | def is_leaf(module): 49 | return sum(1 for x in module.children()) == 0 50 | 51 | 52 | def should_measure(module): 53 | if str(module).startswith("Sequential"): 54 | return False 55 | if is_leaf(module): 56 | return True 57 | return False 58 | 59 | 60 | def cal_multi_adds(model, shape=(2, 3, 32, 32)): 61 | global count_ops 62 | count_ops = 0 63 | data = torch.zeros(shape) 64 | 65 | def new_forward(m): 66 | def lambda_forward(x): 67 | measure_layer(m, x) 68 | return m.old_forward(x) 69 | 70 | return lambda_forward 71 | 72 | def modify_forward(model): 73 | for child in model.children(): 74 | if should_measure(child): 75 | child.old_forward = child.forward 76 | child.forward = new_forward(child) 77 | else: 78 | modify_forward(child) 79 | 80 | def restore_forward(model): 81 | for child in model.children(): 82 | if is_leaf(child) and hasattr(child, "old_forward"): 83 | child.forward = child.old_forward 84 | child.old_forward = None 85 | else: 86 | restore_forward(child) 87 | 88 | modify_forward(model) 89 | model.forward(data) 90 | restore_forward(model) 91 | 92 | return count_ops 93 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from .BaseNormModel import * 2 | from torchvision.models import resnet18, resnet34, resnet50, resnet101 3 | from torchvision.models import alexnet, convnext_tiny, densenet121, efficientnet_b0, googlenet, inception_v3, \ 4 | mnasnet0_75, mobilenet_v3_small, regnet_x_400mf, shufflenet_v2_x0_5, squeezenet1_0, vgg16, \ 5 | vit_b_16, swin_s, maxvit_t, resnet152 6 | from timm.models import adv_inception_v3 7 | from timm.models.inception_resnet_v2 import ens_adv_inception_resnet_v2 8 | from .RobustBench import * 9 | from .SmallResolutionModel import WideResNet_70_16, WideResNet_70_16_dropout 10 | -------------------------------------------------------------------------------- /models/tf_models/DYPmodels.py: -------------------------------------------------------------------------------- 1 | from .tf_inception_v3 import IncV3KitModel 2 | from .tf_inception_v4 import IncV4KitModel 3 | from .tf_inc_res_v2 import IncResV2KitModel 4 | from .tf_ens3_adv_inc_v3 import IncV3Ens3AdvKitModel 5 | from .tf_ens4_adv_inc_v3 import IncV3Ens4AdvKitModel 6 | from .tf_ens_adv_inc_res_v2 import IncResV2EnsKitModel 7 | from .tf_resnet_v2_152 import Res152KitModel 8 | 9 | """ 10 | please download checkpoints from 11 | https://github.com/ylhz/tf_to_pytorch_model 12 | """ 13 | 14 | 15 | def tf_inc_v3(path='./resources/checkpoints/tf_models/tf2torch_inception_v3.npy'): 16 | model = IncV3KitModel(path) 17 | return model 18 | 19 | 20 | def tf_inc_v4(path='./resources/checkpoints/tf_models/tf2torch_inception_v4.npy'): 21 | model = IncV4KitModel(path) 22 | return model 23 | 24 | 25 | def tf_incres_v2(path='./resources/checkpoints/tf_models/tf2torch_inc_res_v2.npy'): 26 | model = IncResV2KitModel(path) 27 | return model 28 | 29 | 30 | def tf_inc_v3_ens3(path='./resources/checkpoints/tf_models/tf2torch_ens3_adv_inc_v3.npy'): 31 | model = IncV3Ens3AdvKitModel(path) 32 | return model 33 | 34 | 35 | def tf_inc_v3_ens4(path='./resources/checkpoints/tf_models/tf2torch_ens4_adv_inc_v3.npy'): 36 | model = IncV3Ens4AdvKitModel(path) 37 | return model 38 | 39 | 40 | def tf_incres_v2_ens(path='./resources/checkpoints/tf_models/tf2torch_ens_adv_inc_res_v2.npy'): 41 | model = IncResV2EnsKitModel(path) 42 | return model 43 | 44 | 45 | def tf_res152(path='./resources/checkpoints/tf_models/tf2torch_resnet_v2_152.npy'): 46 | model = Res152KitModel(path) 47 | return model 48 | -------------------------------------------------------------------------------- /models/tf_models/__init__.py: -------------------------------------------------------------------------------- 1 | from .DYPmodels import tf_inc_v3, tf_inc_v4, tf_incres_v2, \ 2 | tf_inc_v3_ens3, tf_inc_v3_ens4, tf_incres_v2_ens, tf_res152 3 | -------------------------------------------------------------------------------- /models/unets/DiffuserUNets.py: -------------------------------------------------------------------------------- 1 | from diffusers import DDPMPipeline, DDPMScheduler 2 | from torch import nn 3 | 4 | 5 | class DDPMCIFARUnet(nn.Module): 6 | def __init__(self, model_id="google/ddpm-cifar10-32"): 7 | super(DDPMCIFARUnet, self).__init__() 8 | # load model and scheduler 9 | ddpm = DDPMPipeline.from_pretrained( 10 | model_id) # you can replace DDPMPipeline with DDIMPipeline or PNDMPipeline for faster inference 11 | print(ddpm.scheduler) 12 | self.unet = ddpm.unet 13 | 14 | def forward(self, *args, **kwargs): 15 | x = self.unet(*args, **kwargs).sample 16 | return x 17 | 18 | 19 | -------------------------------------------------------------------------------- /models/unets/EDM/__init__.py: -------------------------------------------------------------------------------- 1 | from .edm_nets import EDMPrecond, VPPrecond, VEPrecond 2 | from .get_edm_nets import * 3 | -------------------------------------------------------------------------------- /models/unets/EDM/dnnlib/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is licensed under a Creative Commons 4 | # Attribution-NonCommercial-ShareAlike 4.0 International License. 5 | # You should have received a copy of the license along with this 6 | # work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ 7 | 8 | from .util import EasyDict, make_cache_dir_path 9 | -------------------------------------------------------------------------------- /models/unets/EDM/get_edm_nets.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .edm_nets import EDMPrecond 3 | 4 | 5 | def get_edm_cifar_cond(pretrained=True, **kwargs): 6 | network_kwargs = dict( 7 | model_type="SongUNet", 8 | embedding_type="positional", 9 | encoder_type="standard", 10 | decoder_type="standard", 11 | channel_mult_noise=1, 12 | resample_filter=[1, 1], 13 | model_channels=128, 14 | channel_mult=[2, 2, 2], 15 | augment_dim=9, 16 | ) 17 | network_kwargs.update(kwargs) 18 | model = EDMPrecond(img_resolution=32, img_channels=3, label_dim=10, **network_kwargs) 19 | if pretrained: 20 | model.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_cond.pt")) 21 | return model 22 | 23 | 24 | def get_edm_cifar_uncond(pretrained=True, **kwargs): 25 | network_kwargs = dict( 26 | model_type="SongUNet", 27 | embedding_type="positional", 28 | encoder_type="standard", 29 | decoder_type="standard", 30 | channel_mult_noise=1, 31 | resample_filter=[1, 1], 32 | model_channels=128, 33 | channel_mult=[2, 2, 2], 34 | augment_dim=9, 35 | ) 36 | network_kwargs.update(kwargs) 37 | model = EDMPrecond(img_resolution=32, img_channels=3, label_dim=0, **network_kwargs) 38 | if pretrained: 39 | model.load_state_dict(torch.load("./resources/checkpoints/EDM/edm_cifar_uncond_vp.pt")) 40 | return model 41 | 42 | 43 | def get_edm_cifar_unet(pretrained=True, cond=True, dataset="cifar100", **kwargs): 44 | assert dataset in ["cifar100", "cifar10"] 45 | network_kwargs = dict( 46 | model_type="SongUNet", 47 | embedding_type="positional", 48 | encoder_type="standard", 49 | decoder_type="standard", 50 | channel_mult_noise=1, 51 | resample_filter=[1, 1], 52 | model_channels=128, 53 | channel_mult=[2, 2, 2], 54 | augment_dim=9, 55 | ) 56 | network_kwargs.update(kwargs) 57 | num_classes = int(dataset.split("cifar")[1]) if cond else 0 58 | model = EDMPrecond(img_resolution=32, img_channels=3, label_dim=num_classes, **network_kwargs) 59 | if pretrained: 60 | if dataset == "cifar100": 61 | name = "edm_cifar100_cond_vp.pt" if cond else "edm_cifar100_uncond_vp.pt" 62 | else: 63 | name = "edm_cifar_cond_vp.pt" if cond else "edm_cifar_uncond_vp.pt" 64 | model.load_state_dict(torch.load("./resources/checkpoints/EDM/" + name)) 65 | return model 66 | 67 | 68 | def get_edm_imagenet_64x64_cond(pretrained=True, **kwargs): 69 | network_kwargs = dict(model_type="DhariwalUNet", model_channels=192, channel_mult=[1, 2, 3, 4]) 70 | network_kwargs.update(kwargs) 71 | model = EDMPrecond(img_resolution=64, img_channels=3, label_dim=1000, **network_kwargs) 72 | if pretrained: 73 | model.load_state_dict(torch.load("./resources/checkpoints/EDM/edm-imagenet-64x64-cond.pt")) 74 | return model 75 | -------------------------------------------------------------------------------- /models/unets/EDM/torch_utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is licensed under a Creative Commons 4 | # Attribution-NonCommercial-ShareAlike 4.0 International License. 5 | # You should have received a copy of the license along with this 6 | # work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ 7 | 8 | # empty 9 | -------------------------------------------------------------------------------- /models/unets/EDM/torch_utils/distributed.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is licensed under a Creative Commons 4 | # Attribution-NonCommercial-ShareAlike 4.0 International License. 5 | # You should have received a copy of the license along with this 6 | # work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ 7 | 8 | import os 9 | import torch 10 | from . import training_stats 11 | 12 | #---------------------------------------------------------------------------- 13 | 14 | def init(): 15 | if 'MASTER_ADDR' not in os.environ: 16 | os.environ['MASTER_ADDR'] = 'localhost' 17 | if 'MASTER_PORT' not in os.environ: 18 | os.environ['MASTER_PORT'] = '29500' 19 | if 'RANK' not in os.environ: 20 | os.environ['RANK'] = '0' 21 | if 'LOCAL_RANK' not in os.environ: 22 | os.environ['LOCAL_RANK'] = '0' 23 | if 'WORLD_SIZE' not in os.environ: 24 | os.environ['WORLD_SIZE'] = '1' 25 | 26 | backend = 'gloo' if os.name == 'nt' else 'nccl' 27 | torch.distributed.init_process_group(backend=backend, init_method='env://') 28 | torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', '0'))) 29 | 30 | sync_device = torch.device('cuda') if get_world_size() > 1 else None 31 | training_stats.init_multiprocessing(rank=get_rank(), sync_device=sync_device) 32 | 33 | #---------------------------------------------------------------------------- 34 | 35 | def get_rank(): 36 | return torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 37 | 38 | #---------------------------------------------------------------------------- 39 | 40 | def get_world_size(): 41 | return torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1 42 | 43 | #---------------------------------------------------------------------------- 44 | 45 | def should_stop(): 46 | return False 47 | 48 | #---------------------------------------------------------------------------- 49 | 50 | def update_progress(cur, total): 51 | _ = cur, total 52 | 53 | #---------------------------------------------------------------------------- 54 | 55 | def print0(*args, **kwargs): 56 | if get_rank() == 0: 57 | print(*args, **kwargs) 58 | 59 | #---------------------------------------------------------------------------- 60 | -------------------------------------------------------------------------------- /models/unets/__init__.py: -------------------------------------------------------------------------------- 1 | from .guided_diffusion import get_guided_diffusion_unet 2 | from .score_sde import get_NCSNPP, get_NCSNPP_cached, DiffusionClassifierCached 3 | from .EDM import * 4 | from .DiffuserUNets import * 5 | -------------------------------------------------------------------------------- /models/unets/guided_diffusion/LICENSE_GUIDED_DIFFUSION: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 OpenAI 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /models/unets/guided_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Taken from the following link as is from: 3 | # https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/__init__.py 4 | # 5 | # The license for the original version of this file can be 6 | # found in this directory (LICENSE_GUIDED_DIFFUSION). 7 | # --------------------------------------------------------------- 8 | 9 | """ 10 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 11 | """ 12 | from .default_unet import get_guided_diffusion_unet 13 | from .classifier import create_256x256_classifier 14 | -------------------------------------------------------------------------------- /models/unets/guided_diffusion/classifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .script_util import create_classifier 3 | 4 | 5 | def create_256x256_classifier(pretrained=True): 6 | defaults = dict( 7 | image_size=256, 8 | classifier_use_fp16=False, 9 | classifier_width=128, 10 | classifier_depth=2, 11 | classifier_attention_resolutions="32,16,8", # 16 12 | classifier_use_scale_shift_norm=True, # False 13 | classifier_resblock_updown=True, # False 14 | classifier_pool="attention", 15 | ) 16 | classifier = create_classifier(**defaults) 17 | if pretrained: 18 | classifier.load_state_dict( 19 | torch.load('./resources/checkpoints/guided_diffusion/256x256_classifier.pt') 20 | ) 21 | return classifier 22 | -------------------------------------------------------------------------------- /models/unets/guided_diffusion/config.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | my_config = { 4 | "attention_resolutions": '32,16,8', 5 | "class_cond": False, 6 | "diffusion_steps": 1000, 7 | "rescale_timesteps": True, 8 | "timestep_respacing": '1000', # Modify this value to decrease the number of timesteps. 9 | "image_size": 256, 10 | "learn_sigma": True, 11 | "noise_schedule": 'linear', 12 | "num_channels": 256, 13 | "num_head_channels": 64, 14 | "num_res_blocks": 2, 15 | "resblock_updown": True, 16 | "use_fp16": False, 17 | "use_scale_shift_norm": True, 18 | } 19 | 20 | 21 | 22 | def diffusion_defaults(): 23 | """ 24 | Defaults for image and classifier training. 25 | """ 26 | return dict( 27 | learn_sigma=False, 28 | diffusion_steps=1000, 29 | noise_schedule="linear", 30 | timestep_respacing="", 31 | use_kl=False, 32 | predict_xstart=False, 33 | rescale_timesteps=False, 34 | rescale_learned_sigmas=False, 35 | ) 36 | 37 | 38 | def model_and_diffusion_defaults(): 39 | """ 40 | Defaults for image training. 41 | """ 42 | res = dict( 43 | image_size=64, 44 | num_channels=128, 45 | num_res_blocks=2, 46 | num_heads=4, 47 | num_heads_upsample=-1, 48 | num_head_channels=-1, 49 | attention_resolutions="16,8", 50 | channel_mult="", 51 | dropout=0.0, 52 | class_cond=False, 53 | use_checkpoint=False, 54 | use_scale_shift_norm=True, 55 | resblock_updown=False, 56 | use_fp16=False, 57 | use_new_attention_order=False, 58 | ) 59 | res.update(diffusion_defaults()) 60 | res.update(my_config) 61 | return res 62 | 63 | 64 | -------------------------------------------------------------------------------- /models/unets/guided_diffusion/default_unet.py: -------------------------------------------------------------------------------- 1 | from .script_util import create_model_and_diffusion 2 | from .config import model_and_diffusion_defaults 3 | import torch 4 | 5 | __checkpoint_path__ = "./resources/checkpoints/guided_diffusion/256x256_diffusion" 6 | 7 | 8 | def get_guided_diffusion_unet(pretrained=True, resolution=256, cond=False): 9 | config = model_and_diffusion_defaults() 10 | path = __checkpoint_path__ + "_uncond.pt" if not cond else __checkpoint_path__ + ".pt" 11 | if cond is False: 12 | config["class_cond"] = False 13 | if resolution == 128: 14 | config["image_size"] = 128 15 | elif resolution == 64: 16 | config["image_size"] = 64 17 | model, _ = create_model_and_diffusion(**config) 18 | if pretrained: 19 | model.load_state_dict(torch.load(path.replace("256", str(resolution)))) 20 | 21 | class GuidedDiffusionMeanModel(torch.nn.Module): 22 | def __init__(self): 23 | super(GuidedDiffusionMeanModel, self).__init__() 24 | self.model = model 25 | 26 | def forward(self, *args, **kwargs): 27 | return self.model(*args, **kwargs)[:, :3, :, :] 28 | 29 | return GuidedDiffusionMeanModel() 30 | -------------------------------------------------------------------------------- /models/unets/guided_diffusion/dist_util.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Taken from the following link as is from: 3 | # https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/dist_util.py 4 | # 5 | # The license for the original version of this file can be 6 | # found in this directory (LICENSE_GUIDED_DIFFUSION). 7 | # --------------------------------------------------------------- 8 | 9 | """ 10 | Helpers for distributed training. 11 | """ 12 | 13 | import io 14 | import os 15 | import socket 16 | 17 | import blobfile as bf 18 | from mpi4py import MPI 19 | import torch as th 20 | import torch.distributed as dist 21 | 22 | # Change this to reflect your cluster layout. 23 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 24 | GPUS_PER_NODE = 8 25 | 26 | SETUP_RETRY_COUNT = 3 27 | 28 | 29 | def setup_dist(): 30 | """ 31 | Setup a distributed process group. 32 | """ 33 | if dist.is_initialized(): 34 | return 35 | os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" 36 | 37 | comm = MPI.COMM_WORLD 38 | backend = "gloo" if not th.cuda.is_available() else "nccl" 39 | 40 | if backend == "gloo": 41 | hostname = "localhost" 42 | else: 43 | hostname = socket.gethostbyname(socket.getfqdn()) 44 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 45 | os.environ["RANK"] = str(comm.rank) 46 | os.environ["WORLD_SIZE"] = str(comm.size) 47 | 48 | port = comm.bcast(_find_free_port(), root=0) 49 | os.environ["MASTER_PORT"] = str(port) 50 | dist.init_process_group(backend=backend, init_method="env://") 51 | 52 | 53 | def dev(): 54 | """ 55 | Get the device to use for torch.distributed. 56 | """ 57 | if th.cuda.is_available(): 58 | return th.device(f"cuda") 59 | return th.device("cpu") 60 | 61 | 62 | def load_state_dict(path, **kwargs): 63 | """ 64 | Load a PyTorch file without redundant fetches across MPI ranks. 65 | """ 66 | chunk_size = 2 ** 30 # MPI has a relatively small size limit 67 | if MPI.COMM_WORLD.Get_rank() == 0: 68 | with bf.BlobFile(path, "rb") as f: 69 | data = f.read() 70 | num_chunks = len(data) // chunk_size 71 | if len(data) % chunk_size: 72 | num_chunks += 1 73 | MPI.COMM_WORLD.bcast(num_chunks) 74 | for i in range(0, len(data), chunk_size): 75 | MPI.COMM_WORLD.bcast(data[i : i + chunk_size]) 76 | else: 77 | num_chunks = MPI.COMM_WORLD.bcast(None) 78 | data = bytes() 79 | for _ in range(num_chunks): 80 | data += MPI.COMM_WORLD.bcast(None) 81 | 82 | return th.load(io.BytesIO(data), **kwargs) 83 | 84 | 85 | def sync_params(params): 86 | """ 87 | Synchronize a sequence of Tensors across ranks from rank 0. 88 | """ 89 | for p in params: 90 | with th.no_grad(): 91 | dist.broadcast(p, 0) 92 | 93 | 94 | def _find_free_port(): 95 | try: 96 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 97 | s.bind(("", 0)) 98 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 99 | return s.getsockname()[1] 100 | finally: 101 | s.close() 102 | -------------------------------------------------------------------------------- /models/unets/guided_diffusion/losses.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Taken from the following link as is from: 3 | # https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/losses.py 4 | # 5 | # The license for the original version of this file can be 6 | # found in this directory (LICENSE_GUIDED_DIFFUSION). 7 | # --------------------------------------------------------------- 8 | 9 | """ 10 | Helpers for various likelihood-based losses. These are ported from the original 11 | Ho et al. diffusion models codebase: 12 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 13 | """ 14 | 15 | import numpy as np 16 | 17 | import torch as th 18 | 19 | 20 | def normal_kl(mean1, logvar1, mean2, logvar2): 21 | """ 22 | Compute the KL divergence between two gaussians. 23 | 24 | Shapes are automatically broadcasted, so batches can be compared to 25 | scalars, among other use cases. 26 | """ 27 | tensor = None 28 | for obj in (mean1, logvar1, mean2, logvar2): 29 | if isinstance(obj, th.Tensor): 30 | tensor = obj 31 | break 32 | assert tensor is not None, "at least one argument must be a Tensor" 33 | 34 | # Force variances to be Tensors. Broadcasting helps convert scalars to 35 | # Tensors, but it does not work for th.exp(). 36 | logvar1, logvar2 = [ 37 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 38 | for x in (logvar1, logvar2) 39 | ] 40 | 41 | return 0.5 * ( 42 | -1.0 43 | + logvar2 44 | - logvar1 45 | + th.exp(logvar1 - logvar2) 46 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 47 | ) 48 | 49 | 50 | def approx_standard_normal_cdf(x): 51 | """ 52 | A fast approximation of the cumulative distribution function of the 53 | standard normal. 54 | """ 55 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 56 | 57 | 58 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 59 | """ 60 | Compute the log-likelihood of a Gaussian distribution discretizing to a 61 | given image. 62 | 63 | :param x: the target images. It is assumed that this was uint8 values, 64 | rescaled to the range [-1, 1]. 65 | :param means: the Gaussian mean Tensor. 66 | :param log_scales: the Gaussian log stddev Tensor. 67 | :return: a tensor like x of log probabilities (in nats). 68 | """ 69 | assert x.shape == means.shape == log_scales.shape 70 | centered_x = x - means 71 | inv_stdv = th.exp(-log_scales) 72 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 73 | cdf_plus = approx_standard_normal_cdf(plus_in) 74 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 75 | cdf_min = approx_standard_normal_cdf(min_in) 76 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 77 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 78 | cdf_delta = cdf_plus - cdf_min 79 | log_probs = th.where( 80 | x < -0.999, 81 | log_cdf_plus, 82 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 83 | ) 84 | assert log_probs.shape == x.shape 85 | return log_probs 86 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | # from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = 8 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | 28 | comm = MPI.COMM_WORLD 29 | backend = "gloo" if not th.cuda.is_available() else "nccl" 30 | 31 | if backend == "gloo": 32 | hostname = "localhost" 33 | else: 34 | hostname = socket.gethostbyname(socket.getfqdn()) 35 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 36 | os.environ["RANK"] = str(comm.rank) 37 | os.environ["WORLD_SIZE"] = str(comm.size) 38 | 39 | port = comm.bcast(_find_free_port(), root=0) 40 | os.environ["MASTER_PORT"] = str(port) 41 | dist.init_process_group(backend=backend, init_method="env://") 42 | 43 | 44 | def dev(): 45 | """ 46 | Get the device to use for torch.distributed. 47 | """ 48 | if th.cuda.is_available(): 49 | return th.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}") 50 | return th.device("cpu") 51 | 52 | 53 | def load_state_dict(path, **kwargs): 54 | """ 55 | Load a PyTorch file without redundant fetches across MPI ranks. 56 | """ 57 | if 0 == 0: 58 | with bf.BlobFile(path, "rb") as f: 59 | data = f.read() 60 | else: 61 | data = None 62 | # data = MPI.COMM_WORLD.bcast(data) 63 | return th.load(io.BytesIO(data), **kwargs) 64 | 65 | 66 | def sync_params(params): 67 | """ 68 | Synchronize a sequence of Tensors across ranks from rank 0. 69 | """ 70 | for p in params: 71 | with th.no_grad(): 72 | dist.broadcast(p, 0) 73 | 74 | 75 | def _find_free_port(): 76 | try: 77 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 78 | s.bind(("", 0)) 79 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 80 | return s.getsockname()[1] 81 | finally: 82 | s.close() 83 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import torch.nn as nn 6 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 7 | 8 | 9 | def convert_module_to_f16(l): 10 | """ 11 | Convert primitive modules to float16. 12 | """ 13 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 14 | l.weight.data = l.weight.data.half() 15 | l.bias.data = l.bias.data.half() 16 | 17 | 18 | def convert_module_to_f32(l): 19 | """ 20 | Convert primitive modules to float32, undoing convert_module_to_f16(). 21 | """ 22 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 23 | l.weight.data = l.weight.data.float() 24 | l.bias.data = l.bias.data.float() 25 | 26 | 27 | def make_master_params(model_params): 28 | """ 29 | Copy model parameters into a (differently-shaped) list of full-precision 30 | parameters. 31 | """ 32 | master_params = _flatten_dense_tensors( 33 | [param.detach().float() for param in model_params] 34 | ) 35 | master_params = nn.Parameter(master_params) 36 | master_params.requires_grad = True 37 | return [master_params] 38 | 39 | 40 | def model_grads_to_master_grads(model_params, master_params): 41 | """ 42 | Copy the gradients from the model parameters into the master parameters 43 | from make_master_params(). 44 | """ 45 | master_params[0].grad = _flatten_dense_tensors( 46 | [param.grad.data.detach().float() for param in model_params] 47 | ) 48 | 49 | 50 | def master_params_to_model_params(model_params, master_params): 51 | """ 52 | Copy the master parameter data back into the model parameters. 53 | """ 54 | # Without copying to a list, if a generator is passed, this will 55 | # silently not copy any parameters. 56 | model_params = list(model_params) 57 | 58 | for param, master_param in zip( 59 | model_params, unflatten_master_params(model_params, master_params) 60 | ): 61 | param.detach().copy_(master_param) 62 | 63 | 64 | def unflatten_master_params(model_params, master_params): 65 | """ 66 | Unflatten the master parameters to look like model_params. 67 | """ 68 | return _unflatten_dense_tensors(master_params[0].detach(), model_params) 69 | 70 | 71 | def zero_grad(model_params): 72 | for param in model_params: 73 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 74 | if param.grad is not None: 75 | param.grad.detach_() 76 | param.grad.zero_() 77 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/losses.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import torch as th 4 | 5 | 6 | def normal_kl(mean1, logvar1, mean2, logvar2): 7 | """ 8 | Compute the KL divergence between two gaussians. 9 | 10 | Shapes are automatically broadcasted, so batches can be compared to 11 | scalars, among other use cases. 12 | """ 13 | tensor = None 14 | for obj in (mean1, logvar1, mean2, logvar2): 15 | if isinstance(obj, th.Tensor): 16 | tensor = obj 17 | break 18 | assert tensor is not None, "at least one argument must be a Tensor" 19 | 20 | # Force variances to be Tensors. Broadcasting helps convert scalars to 21 | # Tensors, but it does not work for th.exp(). 22 | logvar1, logvar2 = [ 23 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 24 | for x in (logvar1, logvar2) 25 | ] 26 | 27 | return 0.5 * ( 28 | -1.0 29 | + logvar2 30 | - logvar1 31 | + th.exp(logvar1 - logvar2) 32 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 33 | ) 34 | 35 | 36 | def approx_standard_normal_cdf(x): 37 | """ 38 | A fast approximation of the cumulative distribution function of the 39 | standard normal. 40 | """ 41 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 42 | 43 | 44 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 45 | """ 46 | Compute the log-likelihood of a Gaussian distribution discretizing to a 47 | given image. 48 | 49 | :param x: the target images. It is assumed that this was uint8 values, 50 | rescaled to the range [-1, 1]. 51 | :param means: the Gaussian mean Tensor. 52 | :param log_scales: the Gaussian log stddev Tensor. 53 | :return: a tensor like x of log probabilities (in nats). 54 | """ 55 | assert x.shape == means.shape == log_scales.shape 56 | centered_x = x - means 57 | inv_stdv = th.exp(-log_scales) 58 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 59 | cdf_plus = approx_standard_normal_cdf(plus_in) 60 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 61 | cdf_min = approx_standard_normal_cdf(min_in) 62 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 63 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 64 | cdf_delta = cdf_plus - cdf_min 65 | log_probs = th.where( 66 | x < -0.999, 67 | log_cdf_plus, 68 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 69 | ) 70 | assert log_probs.shape == x.shape 71 | return log_probs 72 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/networks/__init__.py: -------------------------------------------------------------------------------- 1 | from .lenet import * 2 | from .vggnet import * 3 | from .resnet import * 4 | from .wide_resnet import * 5 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/networks/lenet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | 4 | def conv_init(m): 5 | classname = m.__class__.__name__ 6 | if classname.find('Conv') != -1: 7 | init.xavier_uniform(m.weight, gain=np.sqrt(2)) 8 | init.constant(m.bias, 0) 9 | 10 | class LeNet(nn.Module): 11 | def __init__(self, num_classes): 12 | super(LeNet, self).__init__() 13 | self.conv1 = nn.Conv2d(3, 6, 5) 14 | self.conv2 = nn.Conv2d(6, 16, 5) 15 | self.fc1 = nn.Linear(16*5*5, 120) 16 | self.fc2 = nn.Linear(120, 84) 17 | self.fc3 = nn.Linear(84, num_classes) 18 | 19 | def forward(self, x): 20 | out = F.relu(self.conv1(x)) 21 | out = F.max_pool2d(out, 2) 22 | out = F.relu(self.conv2(out)) 23 | out = F.max_pool2d(out, 2) 24 | out = out.view(out.size(0), -1) 25 | out = F.relu(self.fc1(out)) 26 | out = F.relu(self.fc2(out)) 27 | out = self.fc3(out) 28 | 29 | return(out) 30 | -------------------------------------------------------------------------------- /models/unets/improved_diffusion/networks/vggnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | 5 | def conv_init(m): 6 | classname = m.__class__.__name__ 7 | if classname.find('Conv') != -1: 8 | init.xavier_uniform(m.weight, gain=np.sqrt(2)) 9 | init.constant(m.bias, 0) 10 | 11 | def cfg(depth): 12 | depth_lst = [11, 13, 16, 19] 13 | assert (depth in depth_lst), "Error : VGGnet depth should be either 11, 13, 16, 19" 14 | cf_dict = { 15 | '11': [ 16 | 64, 'mp', 17 | 128, 'mp', 18 | 256, 256, 'mp', 19 | 512, 512, 'mp', 20 | 512, 512, 'mp'], 21 | '13': [ 22 | 64, 64, 'mp', 23 | 128, 128, 'mp', 24 | 256, 256, 'mp', 25 | 512, 512, 'mp', 26 | 512, 512, 'mp' 27 | ], 28 | '16': [ 29 | 64, 64, 'mp', 30 | 128, 128, 'mp', 31 | 256, 256, 256, 'mp', 32 | 512, 512, 512, 'mp', 33 | 512, 512, 512, 'mp' 34 | ], 35 | '19': [ 36 | 64, 64, 'mp', 37 | 128, 128, 'mp', 38 | 256, 256, 256, 256, 'mp', 39 | 512, 512, 512, 512, 'mp', 40 | 512, 512, 512, 512, 'mp' 41 | ], 42 | } 43 | 44 | return cf_dict[str(depth)] 45 | 46 | def conv3x3(in_planes, out_planes, stride=1): 47 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=True) 48 | 49 | class VGG(nn.Module): 50 | def __init__(self, depth, num_classes): 51 | super(VGG, self).__init__() 52 | self.features = self._make_layers(cfg(depth)) 53 | self.linear = nn.Linear(512, num_classes) 54 | 55 | def forward(self, x): 56 | out = self.features(x) 57 | out = out.view(out.size(0), -1) 58 | out = self.linear(out) 59 | 60 | return out 61 | 62 | def _make_layers(self, cfg): 63 | layers = [] 64 | in_planes = 3 65 | 66 | for x in cfg: 67 | if x == 'mp': 68 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 69 | else: 70 | layers += [conv3x3(in_planes, x), nn.BatchNorm2d(x), nn.ReLU(inplace=True)] 71 | in_planes = x 72 | 73 | # After cfg convolution 74 | layers += [nn.AvgPool2d(kernel_size=1, stride=1)] 75 | return nn.Sequential(*layers) 76 | 77 | if __name__ == "__main__": 78 | net = VGG(16, 10) 79 | y = net(Variable(torch.randn(1,3,32,32))) 80 | print(y.size()) 81 | -------------------------------------------------------------------------------- /models/unets/score_sde/MultiHeadRDC/DiffusionClassifier.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from .utils import * 4 | 5 | 6 | class DiffusionClassifierCached(nn.Module): 7 | def __init__(self, 8 | unet: nn.Module, 9 | beta: Tensor = torch.linspace(0.1 / 1000, 20 / 1000, 1000), 10 | num_classes=10, 11 | ts: Tensor = torch.arange(1000), 12 | cfg=0, 13 | ): 14 | super(DiffusionClassifierCached, self).__init__() 15 | self.device = torch.device('cuda') 16 | self.unet = unet 17 | alpha = (1 - beta) 18 | self.alpha_bar = alpha.cumprod(dim=0).to(self.device) 19 | self.T = 1000 20 | self.ts = ts.to(self.device) 21 | self._init() 22 | 23 | # storing 24 | self.num_classes = num_classes 25 | self.unet_criterion = nn.MSELoss() 26 | self.cfg = cfg 27 | 28 | def _init(self): 29 | self.eval().requires_grad_(False) 30 | self.to(self.device) 31 | self.transform = lambda x: (x - 0.5) * 2 32 | 33 | def get_one_instance_prediction(self, x: Tensor) -> Tensor: 34 | """ 35 | :param x: 1, C, H, D 36 | """ 37 | loss = self.unet_loss_without_grad(x) 38 | loss = loss * -1 # convert into logit where greatest is the target 39 | return loss 40 | 41 | def forward(self, x: Tensor) -> Tensor: 42 | xs = x.split(1) # 1, C, H, D 43 | y = [] 44 | for now_x in xs: 45 | y.append(self.get_one_instance_prediction(now_x)) 46 | y = torch.stack(y) # N, num_classes 47 | return y 48 | 49 | def unet_loss_without_grad(self, 50 | x: Tensor, 51 | ) -> Tensor: 52 | """ 53 | Calculate the diffusion loss 54 | :param x: in range [0, 1] 55 | """ 56 | t = self.ts 57 | x = x.repeat(t.numel(), 1, 1, 1) 58 | x = self.transform(x) # T, C, H, D 59 | T, C, H, D = x.shape 60 | noise = torch.randn_like(x) 61 | noised_x = torch.sqrt(self.alpha_bar[t]).view(-1, 1, 1, 1) * x + \ 62 | torch.sqrt(1 - self.alpha_bar[t]).view(-1, 1, 1, 1) * noise # T, C, H, D 63 | pre = self.unet(noised_x, t) # T, num_classes*C, H, D 64 | pre = pre.reshape(T, -1, C, H, D).permute(1, 0, 2, 3, 4) # num_classes, T, C, H, D 65 | if self.cfg != 0: 66 | unconditional_prediction = pre[-1:] 67 | conditional_prediction = pre[:-1] 68 | pre = (1 + self.cfg) * conditional_prediction - \ 69 | self.cfg * unconditional_prediction.repeat(self.num_classes, 1, 1, 1, 1) 70 | else: 71 | pre = pre[:self.num_classes] 72 | loss = torch.mean((pre - noise) ** 2, dim=[1, 2, 3, 4]) # num_classes 73 | # pre = pre.reshape(self.num_classes, -1) 74 | # print(cosine_similarity(list(pre.split(1, dim=0)))) 75 | return loss 76 | -------------------------------------------------------------------------------- /models/unets/score_sde/MultiHeadRDC/__init__.py: -------------------------------------------------------------------------------- 1 | from .DiffusionClassifier import DiffusionClassifierCached 2 | -------------------------------------------------------------------------------- /models/unets/score_sde/MultiHeadRDC/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from typing import List 4 | 5 | 6 | def clamp(x, min=0, max=1): 7 | return torch.clamp(x, min=min, max=max) 8 | 9 | 10 | def inplace_clamp(x, min=0, max=1): 11 | return torch.clamp(x, min=min, max=max) 12 | 13 | 14 | def L2Loss(x: Tensor, y: Tensor) -> Tensor: 15 | """ 16 | :param x: N, C, H, D 17 | :param y: N, C, H, D 18 | :return: dim=0 tensor 19 | """ 20 | x = (x - y) ** 2 21 | x = x.view(x.shape[0], -1) 22 | x = torch.norm(x, dim=1, p=2) 23 | x = x.mean(0) 24 | return x 25 | 26 | 27 | def abs_loss(x: Tensor, y: Tensor) -> Tensor: 28 | diff = torch.abs(x - y) 29 | diff = diff.view(diff.shape[0], -1) 30 | diff = torch.sum(diff, dim=1) 31 | diff = torch.mean(diff, dim=0) 32 | return diff 33 | 34 | 35 | def L2_each_instance(x: Tensor, y: Tensor) -> Tensor: 36 | """ 37 | N, ?, ?, ?, ... 38 | :return: N, 39 | """ 40 | x = (x - y) ** 2 41 | x = x.view(x.shape[0], -1) 42 | x = torch.norm(x, dim=1, p=2) 43 | return x 44 | 45 | 46 | def list_mean(x: List) -> float: 47 | return sum(x) / len(x) 48 | -------------------------------------------------------------------------------- /models/unets/score_sde/__init__.py: -------------------------------------------------------------------------------- 1 | from .unets import get_NCSNPP_cached, get_NCSNPP 2 | from .MultiHeadRDC import DiffusionClassifierCached 3 | -------------------------------------------------------------------------------- /models/unets/score_sde/cifar10.yml: -------------------------------------------------------------------------------- 1 | data: 2 | dataset: "CIFAR10" 3 | category: "cifar10" 4 | image_size: 32 5 | num_channels: 3 6 | random_flip: True 7 | centered: True 8 | uniform_dequantization: False 9 | 10 | model: 11 | sigma_min: 0.01 12 | sigma_max: 50 13 | num_scales: 1000 14 | beta_min: 0.1 15 | beta_max: 20. 16 | dropout: 0.1 17 | 18 | name: 'ncsnpp' 19 | scale_by_sigma: False 20 | ema_rate: 0.9999 21 | normalization: 'GroupNorm' 22 | nonlinearity: 'swish' 23 | nf: 128 24 | ch_mult: [1, 2, 2, 2] # (1, 2, 2, 2) 25 | num_res_blocks: 8 26 | attn_resolutions: [16] # (16,) 27 | resamp_with_conv: True 28 | conditional: True 29 | fir: False 30 | fir_kernel: [1, 3, 3, 1] 31 | skip_rescale: True 32 | resblock_type: 'biggan' 33 | progressive: 'none' 34 | progressive_input: 'none' 35 | progressive_combine: 'sum' 36 | attention_type: 'ddpm' 37 | init_scale: 0. 38 | embedding_type: 'positional' 39 | fourier_scale: 16 40 | conv_size: 3 41 | 42 | training: 43 | sde: 'vpsde' 44 | continuous: True 45 | reduce_mean: True 46 | n_iters: 950001 47 | 48 | optim: 49 | weight_decay: 0 50 | optimizer: 'Adam' 51 | lr: 0.0002 # 2e-4 52 | beta1: 0.9 53 | eps: 0.00000001 # 1e-8 54 | warmup: 5000 55 | grad_clip: 1. 56 | 57 | sampling: 58 | n_steps_each: 1 59 | noise_removal: True 60 | probability_flow: False 61 | snr: 0.16 62 | 63 | method: 'pc' 64 | predictor: 'euler_maruyama' 65 | corrector: 'none' -------------------------------------------------------------------------------- /models/unets/score_sde/models/__init__.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | # Copyright 2020 The Google Research Authors. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | from . import ncsnpp 16 | -------------------------------------------------------------------------------- /models/unets/score_sde/op/__init__.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Taken from the following link as is from: 3 | # https://github.com/yang-song/score_sde_pytorch/blob/main/op/__init__.py 4 | # 5 | # The license for the original version of this file can be 6 | # found in the `score_sde` directory (LICENSE_SCORE_SDE). 7 | # --------------------------------------------------------------- 8 | 9 | from .fused_act import FusedLeakyReLU, fused_leaky_relu 10 | from .upfirdn2d import upfirdn2d 11 | -------------------------------------------------------------------------------- /models/unets/score_sde/op/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Taken from the following link as is from: 3 | // https://github.com/yang-song/score_sde_pytorch/blob/main/op/fused_bias_act.cpp 4 | // 5 | // The license for the original version of this file can be 6 | // found in the `score_sde` directory (LICENSE_SCORE_SDE). 7 | // --------------------------------------------------------------- 8 | 9 | #include 10 | 11 | 12 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 13 | int act, int grad, float alpha, float scale); 14 | 15 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 16 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 17 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 18 | 19 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 20 | int act, int grad, float alpha, float scale) { 21 | CHECK_CUDA(input); 22 | CHECK_CUDA(bias); 23 | 24 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 25 | } 26 | 27 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 28 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 29 | } -------------------------------------------------------------------------------- /models/unets/score_sde/op/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | // 3 | // This work is made available under the Nvidia Source Code License-NC. 4 | // To view a copy of this license, visit 5 | // https://nvlabs.github.io/stylegan2/license.html 6 | 7 | #include 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | 17 | 18 | template 19 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 20 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 21 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 22 | 23 | scalar_t zero = 0.0; 24 | 25 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 26 | scalar_t x = p_x[xi]; 27 | 28 | if (use_bias) { 29 | x += p_b[(xi / step_b) % size_b]; 30 | } 31 | 32 | scalar_t ref = use_ref ? p_ref[xi] : zero; 33 | 34 | scalar_t y; 35 | 36 | switch (act * 10 + grad) { 37 | default: 38 | case 10: y = x; break; 39 | case 11: y = x; break; 40 | case 12: y = 0.0; break; 41 | 42 | case 30: y = (x > 0.0) ? x : x * alpha; break; 43 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 44 | case 32: y = 0.0; break; 45 | } 46 | 47 | out[xi] = y * scale; 48 | } 49 | } 50 | 51 | 52 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 53 | int act, int grad, float alpha, float scale) { 54 | int curDevice = -1; 55 | cudaGetDevice(&curDevice); 56 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 57 | 58 | auto x = input.contiguous(); 59 | auto b = bias.contiguous(); 60 | auto ref = refer.contiguous(); 61 | 62 | int use_bias = b.numel() ? 1 : 0; 63 | int use_ref = ref.numel() ? 1 : 0; 64 | 65 | int size_x = x.numel(); 66 | int size_b = b.numel(); 67 | int step_b = 1; 68 | 69 | for (int i = 1 + 1; i < x.dim(); i++) { 70 | step_b *= x.size(i); 71 | } 72 | 73 | int loop_x = 4; 74 | int block_size = 4 * 32; 75 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 76 | 77 | auto y = torch::empty_like(x); 78 | 79 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 80 | fused_bias_act_kernel<<>>( 81 | y.data_ptr(), 82 | x.data_ptr(), 83 | b.data_ptr(), 84 | ref.data_ptr(), 85 | act, 86 | grad, 87 | alpha, 88 | scale, 89 | loop_x, 90 | size_x, 91 | step_b, 92 | size_b, 93 | use_bias, 94 | use_ref 95 | ); 96 | }); 97 | 98 | return y; 99 | } -------------------------------------------------------------------------------- /models/unets/score_sde/op/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Taken from the following link as is from: 3 | // https://github.com/yang-song/score_sde_pytorch/blob/main/op/upfirdn2d.cpp 4 | // 5 | // The license for the original version of this file can be 6 | // found in the `score_sde` directory (LICENSE_SCORE_SDE). 7 | // --------------------------------------------------------------- 8 | 9 | #include 10 | 11 | 12 | torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, 13 | int up_x, int up_y, int down_x, int down_y, 14 | int pad_x0, int pad_x1, int pad_y0, int pad_y1); 15 | 16 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 17 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 18 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 19 | 20 | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, 21 | int up_x, int up_y, int down_x, int down_y, 22 | int pad_x0, int pad_x1, int pad_y0, int pad_y1) { 23 | CHECK_CUDA(input); 24 | CHECK_CUDA(kernel); 25 | 26 | return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); 27 | } 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); 31 | } -------------------------------------------------------------------------------- /optimizer/ALRS.py: -------------------------------------------------------------------------------- 1 | class IdentityScheduler: 2 | def __init__(self, *args, **kwargs): 3 | pass 4 | 5 | def step(self, *args, **kwargs): 6 | pass 7 | 8 | 9 | class ALRS: 10 | """ 11 | refer to 12 | Bootstrap Generalization Ability from Loss Landscape Perspective 13 | """ 14 | 15 | def __init__(self, optimizer, loss_threshold=0.01, loss_ratio_threshold=0.01, decay_rate=0.97): 16 | self.optimizer = optimizer 17 | self.loss_threshold = loss_threshold 18 | self.decay_rate = decay_rate 19 | self.loss_ratio_threshold = loss_ratio_threshold 20 | 21 | self.last_loss = 999 22 | 23 | def step(self, loss): 24 | delta = self.last_loss - loss 25 | if delta < self.loss_threshold and delta / self.last_loss < self.loss_ratio_threshold: 26 | for group in self.optimizer.param_groups: 27 | group["lr"] *= self.decay_rate 28 | now_lr = group["lr"] 29 | print(f"now lr = {now_lr}") 30 | 31 | self.last_loss = loss 32 | -------------------------------------------------------------------------------- /optimizer/PGD.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.optim import Optimizer 3 | 4 | 5 | class PGD(Optimizer): 6 | def __init__(self, params, lr=5e-3, maximum=True, epsilon=16 / 255): 7 | dampening = 0 8 | weight_decay = 0 9 | nesterov = False 10 | maximize = False 11 | momentum = 0 12 | defaults = dict(lr=lr, momentum=momentum, dampening=dampening, 13 | weight_decay=weight_decay, nesterov=nesterov, maximize=maximize) 14 | super(PGD, self).__init__(params, defaults) 15 | self.maximum = maximum 16 | self.lr = lr 17 | self.original_params = [p.clone() for p in params] 18 | self.epsilon = epsilon 19 | 20 | @torch.no_grad() 21 | def step(self, closure=None): 22 | if self.maximum: 23 | for group in self.param_groups: 24 | for i, p in enumerate(group['params']): 25 | if p.grad is not None: 26 | p.add_(self.lr * p.grad.sign()) 27 | p.clamp_(min=self.original_params[i] - self.epsilon, 28 | max=self.original_params[i] + self.epsilon) 29 | else: 30 | for group in self.param_groups: 31 | for i, p in enumerate(group['params']): 32 | if p.grad is not None: 33 | p.add_(-self.lr * p.grad.sign()) 34 | p.clamp_(min=self.original_params[i] - self.epsilon, 35 | max=self.original_params[i] + self.epsilon) 36 | -------------------------------------------------------------------------------- /optimizer/__init__.py: -------------------------------------------------------------------------------- 1 | from .PGD import PGD 2 | from torch.optim import Adam, AdamW, SGD, Optimizer 3 | from .default import default_optimizer, default_lr_scheduler 4 | from .ALRS import * 5 | -------------------------------------------------------------------------------- /optimizer/default.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from .ALRS import ALRS 4 | 5 | 6 | def default_optimizer(model: nn.Module, lr=1e-1, ) -> torch.optim.Optimizer: 7 | return torch.optim.SGD(model.parameters(), lr=lr) 8 | 9 | 10 | def default_lr_scheduler(optimizer): 11 | return ALRS(optimizer) 12 | -------------------------------------------------------------------------------- /optimizer/losses/CKA.py: -------------------------------------------------------------------------------- 1 | """ 2 | Similarity of Neural Network Representations Revisited 3 | """ 4 | from torch import nn, Tensor 5 | import torch 6 | 7 | 8 | class UnNormalizedSimpleCKA(nn.Module): 9 | def __init__(self): 10 | super(UnNormalizedSimpleCKA, self).__init__() 11 | 12 | def forward(self, x: Tensor, y: Tensor) -> Tensor: 13 | """ 14 | x, y should be (N, D) tensor 15 | """ 16 | assert x.shape == y.shape, 'feature shape should be equal if you use CKA' 17 | return torch.sum((x.T @ y) ** 2) 18 | 19 | 20 | class NormalizedSimpleCKA(nn.Module): 21 | def __init__(self): 22 | super(NormalizedSimpleCKA, self).__init__() 23 | 24 | def forward(self, x: Tensor, y: Tensor) -> Tensor: 25 | """ 26 | x, y should be (N, D) tensor 27 | """ 28 | assert x.shape == y.shape, 'feature shape should be equal if you use CKA' 29 | x_f, y_f = torch.norm(x @ x.T, p='fro'), torch.norm(y @ y.T, p='fro') 30 | return torch.sum((x.T @ y) ** 2) / (x_f * y_f) 31 | 32 | 33 | class CenteredKernelAnalysis(nn.Module): 34 | def __init__(self, kernel_function): 35 | """ 36 | :param kernel_function: given a matrix (N, D), return a matrix K with shape (N, N), 37 | where K_{ij} = kernel(x_i, x_j) 38 | """ 39 | super(CenteredKernelAnalysis, self).__init__() 40 | self.kernel_function = kernel_function 41 | 42 | def forward(self, x: Tensor, y: Tensor) -> Tensor: 43 | """ 44 | x, y should be (N, D) tensor 45 | """ 46 | assert x.shape == y.shape, 'feature shape should be equal if you use CKA' 47 | k_x, k_y = self.kernel_function(x), self.kernel_function(y) 48 | x_f, y_f = torch.norm(k_x, p='fro'), torch.norm(k_y, p='fro') 49 | return (k_x.view(-1) @ k_y.view(-1)) / (x_f * y_f) 50 | -------------------------------------------------------------------------------- /optimizer/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from .CKA import * 2 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==1.4.0 2 | accelerate==0.19.0 3 | autoattack==0.1 4 | bleach==6.0.0 5 | blessed==1.20.0 6 | boltons==23.0.0 7 | brotlipy==0.7.0 8 | cachetools==5.3.0 9 | certifi==2022.12.7 10 | cffi==1.15.1 11 | chardet==3.0.4 12 | charset-normalizer==2.0.4 13 | cmake==3.26.0 14 | conda==23.1.0 15 | conda-content-trust==0.1.3 16 | conda-package-handling==2.0.2 17 | conda_package_streaming==0.7.0 18 | contextlib2==21.6.0 19 | contourpy==1.0.7 20 | cryptography==38.0.4 21 | cycler==0.11.0 22 | diffusers~=0.18.2 23 | einops==0.6.1 24 | filelock==3.10.0 25 | fonttools==4.39.2 26 | fsspec==2023.5.0 27 | functorch==2.0.0 28 | geotorch==0.3.0 29 | google-auth==2.16.2 30 | google-auth-oauthlib==0.4.6 31 | gpustat==1.1 32 | grpcio==1.51.3 33 | huggingface-hub==0.14.1 34 | idna==2.10 35 | importlib-metadata==6.6.0 36 | Jinja2==3.1.2 37 | kaggle==1.5.15 38 | kiwisolver==1.4.4 39 | lit==15.0.7 40 | Markdown==3.4.1 41 | MarkupSafe==2.1.2 42 | matplotlib~=3.5.1 43 | ml-collections==0.1.1 44 | mpmath==1.3.0 45 | networkx==3.0 46 | ninja==1.11.1 47 | numpy~=1.25.2 48 | nvidia-cublas-cu11==11.10.3.66 49 | nvidia-cuda-cupti-cu11==11.7.101 50 | nvidia-cuda-nvrtc-cu11==11.7.99 51 | nvidia-cuda-runtime-cu11==11.7.99 52 | nvidia-cudnn-cu11==8.5.0.96 53 | nvidia-cufft-cu11==10.9.0.58 54 | nvidia-curand-cu11==10.2.10.91 55 | nvidia-cusolver-cu11==11.4.0.1 56 | nvidia-cusparse-cu11==11.7.4.91 57 | nvidia-ml-py==11.525.112 58 | nvidia-nccl-cu11==2.14.3 59 | nvidia-nvtx-cu11==11.7.91 60 | oauthlib==3.2.2 61 | opencv-python~=4.8.0.76 62 | packaging==23.0 63 | pandas==1.3.5 64 | patsy==0.5.3 65 | Pillow~=9.1.0 66 | pip==22.3.1 67 | pluggy==1.0.0 68 | protobuf==4.22.1 69 | psutil==5.9.5 70 | pyasn1==0.4.8 71 | pyasn1-modules==0.2.8 72 | pycosat==0.6.4 73 | pycparser==2.21 74 | pyOpenSSL==22.0.0 75 | pyparsing==3.0.9 76 | PySocks==1.7.1 77 | python-dateutil==2.8.2 78 | python-slugify==8.0.1 79 | pytorch-fid==0.3.0 80 | pytz==2022.7.1 81 | PyYAML~=5.4.1 82 | regex==2023.5.5 83 | requests~=2.25.1 84 | requests-oauthlib==1.3.1 85 | robustbench==1.1 86 | rsa==4.9 87 | ruamel.yaml==0.17.21 88 | ruamel.yaml.clib==0.2.6 89 | scipy~=1.8.1 90 | sentencepiece==0.1.99 91 | setuptools==65.6.3 92 | six==1.16.0 93 | statsmodels~=0.13.2 94 | sympy==1.11.1 95 | tensorboard==2.12.0 96 | tensorboard-data-server==0.7.0 97 | tensorboard-plugin-wit==1.8.1 98 | text-unidecode==1.3 99 | timm==0.6.12 100 | tokenizers==0.13.3 101 | toolz==0.12.0 102 | torch==2.0.0 103 | torchaudio==2.0.1 104 | torchdiffeq==0.2.3 105 | torchode==0.1.8 106 | torchsde==0.2.5 107 | torchtyping==0.1.4 108 | torchvision==0.15.1 109 | tqdm~=4.64.0 110 | trampoline==0.1.2 111 | transformers~=4.31.0 112 | triton==2.0.0 113 | typeguard==3.0.2 114 | typing_extensions==4.5.0 115 | urllib3==1.26.14 116 | wcwidth==0.2.6 117 | webencodings==0.5.1 118 | Werkzeug==2.2.3 119 | wheel==0.37.1 120 | zipp==3.15.0 121 | zstandard==0.18.0 122 | torchopt 123 | yaml~=0.2.5 -------------------------------------------------------------------------------- /tester/AutoAttack.py: -------------------------------------------------------------------------------- 1 | # from autoattack import AutoAttack 2 | from attacks.autoattack import AutoAttack 3 | import torch 4 | from torch import nn, Tensor 5 | from torch.utils.data import DataLoader 6 | 7 | 8 | def test_autoattack_acc(model: nn.Module, loader: DataLoader, bs=16, log_path=None): 9 | adversary = AutoAttack(model, eps=8 / 255, log_path=log_path) 10 | xs, ys = [], [] 11 | for x, y in loader: 12 | xs.append(x) 13 | ys.append(y) 14 | x = torch.concat(xs, dim=0).cuda() 15 | y = torch.concat(ys, dim=0).cuda() 16 | adversary.run_standard_evaluation(x, y, bs=bs) 17 | 18 | 19 | def test_apgd_dlr_acc(model: nn.Module, 20 | x: Tensor = None, 21 | y: Tensor = None, 22 | loader: DataLoader = None, 23 | bs=1, 24 | log_path=None, device=torch.device('cuda'), 25 | eps=8 / 255, 26 | norm='Linf', 27 | eot_iter=1, 28 | n_iter=100, 29 | ): 30 | model.eval().requires_grad_(False).cuda() 31 | if loader is not None: 32 | xs, ys = [], [] 33 | for x, y in loader: 34 | xs.append(x) 35 | ys.append(y) 36 | x = torch.concat(xs, dim=0).cuda() 37 | y = torch.concat(ys, dim=0).cuda() 38 | adversary = AutoAttack(model, norm=norm, eps=eps, version='custom', 39 | log_path=log_path, device=device) 40 | adversary.attacks_to_run = ['apgd-dlr'] 41 | adversary.apgd.eot_iter = eot_iter 42 | adversary.apgd.n_iter = n_iter 43 | adversary.run_standard_evaluation(x, y, bs=bs) 44 | 45 | ## 46 | # 47 | # 48 | # 49 | # 50 | # 51 | # def test_apgd_dlr_multi_gpu(model: nn.Module, loader: DataLoader, bs=1, 52 | # ngpus=4): 53 | # xs, ys = [], [] 54 | # for x, y in loader: 55 | # xs.append(x) 56 | # ys.append(y) 57 | # x = torch.concat(xs, dim=0) 58 | # y = torch.concat(ys, dim=0) 59 | # x = x.split(math.floor(x.shape[0] / ngpus)) 60 | # y = y.split(math.floor(y.shape[0] / ngpus)) 61 | # processes = [] 62 | # for i in range(ngpus): 63 | # now_device = torch.device(f'cuda:{i}') 64 | # now = multiprocessing.Process(target=test_apgd_dlr_acc, 65 | # args=( 66 | # model.to(now_device), 67 | # x[i].to(now_device), 68 | # y[i].to(now_device), 69 | # bs, 70 | # f'apgd_device_{i}.txt', 71 | # now_device, 72 | # ) 73 | # ) 74 | # processes.append(now) 75 | # for now in processes: 76 | # now.start() 77 | # for now in processes: 78 | # now.join() 79 | -------------------------------------------------------------------------------- /tester/CertifyRobustness.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | from typing import Iterable, Dict, Tuple, List 4 | from tqdm import tqdm 5 | from defenses.RandomizedSmoothing import Smooth 6 | from scipy.stats import norm 7 | from statsmodels.stats.proportion import proportion_confint 8 | 9 | __WRONG_PREDICTION__ = -1 10 | 11 | 12 | def robust_radius_given_correct_num(nA: int = 990, n: int = 1000, alpha: float = 0.001, sigma: float = 0.25) -> float: 13 | pABar = proportion_confint(nA, n, alpha=2 * alpha, method="beta")[0] 14 | radius = sigma * norm.ppf(pABar) 15 | return radius 16 | 17 | 18 | def radii_discretion(radii: List[float], epsilons: Iterable = (0, 0.25, 0.5, 0.75, 1)): 19 | radii_tensor = torch.tensor(radii) 20 | denominator = len(radii) 21 | result = dict() 22 | for eps in epsilons: 23 | print("-" * 100) 24 | result[eps] = torch.sum(radii_tensor >= eps).item() / denominator 25 | print(f"certified robustness at {eps} is {result[eps]}") 26 | print("-" * 100) 27 | return result 28 | 29 | 30 | def nA_and_n_to_radii(nAs: List, ns: List, *args, **kwargs): 31 | radii = [] 32 | for nA, n in zip(nAs, ns): 33 | radii.append(robust_radius_given_correct_num(nA, n, *args, **kwargs)) 34 | return radii 35 | 36 | 37 | @torch.no_grad() 38 | def certify_robustness( 39 | model: Smooth, 40 | loader: DataLoader or Iterable, 41 | epsilons: Iterable = (0, 0.25, 0.5, 0.75, 1), 42 | device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), 43 | *args, 44 | **kwargs, 45 | ) -> Tuple[Dict, List, List, List]: 46 | model.base_classifier.to(device).eval() 47 | radii, nAs, ns = [], [], [] 48 | for x, y in tqdm(loader): 49 | x, y = x.to(device), y.to(device) 50 | label, radius, nA, n = model.certify(x.squeeze(), *args, **kwargs) 51 | radii.append(radius if label == y.item() else __WRONG_PREDICTION__) 52 | nAs.append(nA if label == y.item() else 0) 53 | ns.append(n) 54 | result = radii_discretion(radii, epsilons) 55 | return result, nAs, ns, radii 56 | -------------------------------------------------------------------------------- /tester/__init__.py: -------------------------------------------------------------------------------- 1 | from .TestAcc import * 2 | from .TransferAttackAcc import test_transfer_attack_acc, test_robustness 3 | from .AutoAttack import * 4 | from .CertifyRobustness import certify_robustness 5 | -------------------------------------------------------------------------------- /tester/generation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | import os 4 | from typing import Callable, Iterable 5 | from torchvision import transforms 6 | from tqdm import tqdm 7 | 8 | try: 9 | import pytorch_fid 10 | except: 11 | assert False, 'You should install pytorch_fid before testing fid' 12 | 13 | 14 | @torch.no_grad() 15 | def random_label_generator(num_classes=10, 16 | batch_size=64, 17 | device=torch.device('cuda')): 18 | return torch.randint(low=0, high=num_classes, size=(batch_size,), device=device) 19 | 20 | 21 | @torch.no_grad() 22 | def test_fid(loader: DataLoader or Iterable, generator: Callable, 23 | path: str = './test_fid/', 24 | dataset_sample_images: int = 50000, 25 | total_generation_iter: int = 100, 26 | ): 27 | to_img = transforms.ToPILImage() 28 | if not os.path.exists(path): 29 | os.mkdir(path) 30 | dataset_path = os.path.join(path, 'dataset/') 31 | if not os.path.exists(dataset_path): 32 | os.mkdir(dataset_path) 33 | generated_img_path = os.path.join(path, 'generated/') 34 | if not os.path.exists(generated_img_path): 35 | os.mkdir(generated_img_path) 36 | 37 | # extract images from loader 38 | imgs = [] 39 | for x, _ in loader: 40 | imgs += list(torch.split(x, split_size_or_sections=1, dim=0)) 41 | imgs = imgs[:dataset_sample_images] 42 | for i, img in enumerate(imgs): 43 | img = to_img(img.squeeze()) 44 | img.save(os.path.join(dataset_path, f'{i}.png')) 45 | # extract images from generator 46 | imgs.clear() 47 | for _ in tqdm(range(total_generation_iter)): 48 | imgs += list(torch.split(generator(), split_size_or_sections=1, dim=0)) 49 | for i, img in enumerate(imgs): 50 | img = to_img(img.squeeze()) 51 | img.save(os.path.join(generated_img_path, f'{i}.png')) 52 | 53 | # test fid 54 | os.system(f'python -m pytorch_fid {dataset_path} {generated_img_path}') 55 | -------------------------------------------------------------------------------- /tester/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os 3 | from typing import List 4 | 5 | 6 | def cosine_similarity(x: List): 7 | ''' 8 | input a list of tensor with same shape. return the mean cosine_similarity 9 | ''' 10 | x = torch.stack(x) 11 | N = x.shape[0] 12 | x = x.reshape(N, -1) 13 | 14 | norm = torch.norm(x, p=2, dim=1) 15 | x /= norm.reshape(-1, 1) # N, D 16 | similarity = x @ x.T # N, N 17 | # only the tri-upper part. Note that the diagonal means which one is diagonal. 18 | mask = torch.triu(torch.ones(N, N, device=x.device), diagonal=1).to(torch.bool) 19 | similarity = similarity[mask] 20 | return torch.mean(similarity).item() 21 | 22 | 23 | def list_mean(x: list) -> float: 24 | return sum(x) / len(x) 25 | 26 | 27 | def update_log(x: str or List[str], path: str = './log.txt'): 28 | with open(path, 'a') as file: 29 | if type(x) is str: 30 | file.write(x if x.endswith('\n') else x + '\n') 31 | elif type(x) is list: 32 | for now in x: 33 | file.write(now if now.endswith('\n') else now + '\n') 34 | -------------------------------------------------------------------------------- /utils/HRNet/Changer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch import fx 4 | from typing import Dict, Any, Tuple 5 | from .HRBatchNorm import HRBatchNorm2D 6 | from .HRReLU import HRReLU 7 | from .HRLinear import HRLinear 8 | 9 | __all__ = ['change'] 10 | 11 | 12 | def _parent_name(target: str) -> Tuple[str, str]: 13 | """ 14 | Splits a qualname into parent path and last atom. 15 | For example, `foo.bar.baz` -> (`foo.bar`, `baz`) 16 | """ 17 | *parent, name = target.rsplit('.', 1) 18 | return parent[0] if parent else '', name 19 | 20 | 21 | def replace_node_module(node: fx.Node, modules: Dict[str, Any], new_module: torch.nn.Module): 22 | assert (isinstance(node.target, str)) 23 | parent_name, name = _parent_name(node.target) 24 | modules[node.target] = new_module 25 | setattr(modules[parent_name], name, new_module) 26 | 27 | 28 | @torch.no_grad() 29 | def change(model: nn.Module) -> fx.GraphModule: 30 | traced_model = fx.symbolic_trace(model) 31 | named_modules = dict(traced_model.named_modules()) 32 | for node in traced_model.graph.nodes: 33 | if node.op == 'call_module': 34 | now_module = named_modules[node.target] 35 | if isinstance(now_module, nn.Linear): 36 | weight = now_module.weight 37 | bias = now_module.bias 38 | new_module = HRLinear(weight, bias) 39 | replace_node_module(node, named_modules, new_module) 40 | if isinstance(now_module, nn.ReLU): 41 | replace_node_module(node, named_modules, HRReLU()) 42 | if isinstance(now_module, nn.BatchNorm2d): 43 | weight = now_module.weight 44 | bias = now_module.bias 45 | mean = now_module.running_mean 46 | std = torch.sqrt(now_module.running_var + now_module.eps) 47 | new_module = HRBatchNorm2D(mean, std, weight, bias) 48 | replace_node_module(node, named_modules, new_module) 49 | return traced_model 50 | -------------------------------------------------------------------------------- /utils/HRNet/HRBatchNorm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | from torch.autograd import Function 4 | from typing import Any, Optional 5 | from torch.nn import functional as F 6 | 7 | 8 | class HRBatchNorm2DFunction(Function): 9 | @staticmethod 10 | def forward(ctx: Any, x: Tensor, weight: Tensor, bias: Tensor) -> Any: 11 | """ 12 | :param ctx: 13 | :param x: N, C, H, D 14 | :param weight: C 15 | :param bias: C 16 | :return: 17 | """ 18 | ctx.weight = weight 19 | x = x.permute(0, 2, 3, 1) # N, H, D, C 20 | x = x * weight + bias 21 | x = x.permute(0, 3, 1, 2) # N, C, H, D 22 | return x 23 | 24 | @staticmethod 25 | def backward(ctx: Any, dy) -> Any: 26 | """ 27 | :param ctx: 28 | :param dy: N, C, H, D 29 | :return: 30 | """ 31 | dy = dy.permute(0, 2, 3, 1) # N, H, D, C 32 | weight = ctx.weight 33 | dx = dy * weight 34 | dx = dx.permute(0, 3, 1, 2) 35 | return dx, None, None 36 | 37 | 38 | class HRBatchNorm2D(nn.Module): 39 | def __init__(self, mean: Tensor, std: Tensor, weight: Tensor, bias: Tensor): 40 | """ 41 | y = weight*(x-mean)/std+bias = weight/std * x + bias - weight/std * mean 42 | :param mean:C 43 | :param std:C 44 | :param weight:C 45 | :param bias:C 46 | """ 47 | super(HRBatchNorm2D, self).__init__() 48 | self.weight = weight / std 49 | self.bias = bias - self.weight * mean 50 | 51 | def forward(self, x: Tensor): 52 | """ 53 | :param x: N, C, H, D 54 | """ 55 | return HRBatchNorm2DFunction.apply(x, self.weight, self.bias) 56 | -------------------------------------------------------------------------------- /utils/HRNet/HRConv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | from torch.autograd import Function 4 | from typing import Any, Optional 5 | from torch.nn import functional as F 6 | import torchvision 7 | 8 | 9 | """ 10 | 11 | """ 12 | 13 | 14 | class HRConvFunction(Function): 15 | @staticmethod 16 | def forward(ctx: Any, input: Tensor, weight: Tensor, bias: Optional[Tensor], 17 | stride: int, padding: int = 0, dilation: bool = False, groups: int = 1) -> Any: 18 | ctx.weight = weight 19 | return F.conv2d(input, weight, bias, stride, 20 | padding, dilation, groups) 21 | 22 | @staticmethod 23 | def backward(ctx: Any, dy) -> Any: 24 | weight = ctx.weight 25 | 26 | 27 | class HRConv(nn.Module): 28 | """ 29 | padding only support zero padding!! 30 | not support dilation convolution!!!!!! 31 | 32 | """ 33 | 34 | def __init__(self, weight: Tensor, bias: Optional[Tensor], 35 | stride: int, padding: int = 0, dilation: bool = False, groups: int = 1): 36 | super(HRConv, self).__init__() 37 | 38 | def forward(self, input: Tensor): 39 | pass -------------------------------------------------------------------------------- /utils/HRNet/HRLinear.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | from torch.autograd import Function 4 | from typing import Any 5 | import torch.nn.functional as F 6 | 7 | 8 | class HRLinearFunction(Function): 9 | @staticmethod 10 | def forward(ctx: Any, x: Tensor, weight: Tensor, bias: Tensor) -> Any: 11 | """ 12 | :param x: N, D1 -> N, D2 13 | :param weight: D2, D1 14 | :param bias: D2 15 | :return: 16 | """ 17 | ctx.weight = weight 18 | return F.linear(x, weight, bias) 19 | 20 | @staticmethod 21 | def backward(ctx: Any, dy: Tensor) -> Any: 22 | """ 23 | :param dy: N, D2 24 | :return: 25 | """ 26 | weight = ctx.weight # D2, D1 27 | dx = dy @ weight # N, D1 28 | return dx, None, None 29 | 30 | 31 | class HRLinear(nn.Module): 32 | def __init__(self, weight: Tensor, bias: Tensor): 33 | super(HRLinear, self).__init__() 34 | self.weight = weight 35 | self.bias = bias 36 | # self.weight.requires_grad_(False) 37 | # self.bias.requires_grad_(False) 38 | 39 | def forward(self, x: Tensor) -> Tensor: 40 | return HRLinearFunction.apply(x, self.weight, self.bias) 41 | -------------------------------------------------------------------------------- /utils/HRNet/HRReLU.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn, Tensor 3 | from torch.autograd import Function 4 | from typing import Any 5 | 6 | 7 | class HRReLUFunction(Function): 8 | @staticmethod 9 | def forward(ctx: Any, x: Tensor) -> Tensor: 10 | mask = x < 0 11 | ctx.mask = mask 12 | x[mask] = 0 13 | return x 14 | 15 | @staticmethod 16 | def backward(ctx: Any, dy: Tensor) -> Tensor: 17 | mask = ctx.mask 18 | grad = dy.clone() 19 | grad[mask] = 0 20 | return grad 21 | 22 | 23 | class HRReLU(nn.Module): 24 | def __init__(self): 25 | super(HRReLU, self).__init__() 26 | 27 | def forward(self, x: Tensor): 28 | return HRReLUFunction.apply(x) 29 | -------------------------------------------------------------------------------- /utils/HRNet/__init__.py: -------------------------------------------------------------------------------- 1 | from .Changer import change -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .ImageHandling import * 2 | from .plot import * 3 | from .HRNet import * 4 | from .gpu import * 5 | from .seed import * 6 | -------------------------------------------------------------------------------- /utils/gpu.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import time 3 | 4 | 5 | def hold_gpu(memory: float = 24): 6 | size = 655360000 * 9 * memory / 24 7 | x = torch.randn(int(size), device=torch.device('cuda')) 8 | time.sleep(86400 * 4) 9 | -------------------------------------------------------------------------------- /utils/plot/ColorUtils.py: -------------------------------------------------------------------------------- 1 | import random 2 | from matplotlib.colors import ListedColormap 3 | import datetime 4 | import os 5 | 6 | white = (1, 1, 1) 7 | 8 | """ 9 | 10 | """ 11 | 12 | 13 | def rand_color() -> tuple: 14 | return (random.random(), random.random(), random.random()) 15 | 16 | 17 | def get_rand_cmap(): 18 | return ListedColormap((white, rand_color())) 19 | 20 | 21 | def get_datetime_str(style='dt'): 22 | cur_time = datetime.datetime.now() 23 | date_str = cur_time.strftime('%y_%m_%d_') 24 | time_str = cur_time.strftime('%H_%M_%S') 25 | if style == 'data': 26 | return date_str 27 | elif style == 'time': 28 | return time_str 29 | return date_str + time_str 30 | 31 | 32 | class suppress_stdout_stderr(object): 33 | ''' 34 | A context manager for doing a "deep suppression" of stdout and stderr in 35 | Python, i.e. will suppress all print, even if the print originates in a 36 | compiled C/Fortran sub-function. 37 | This will not suppress raised exceptions, since exceptions are printed 38 | to stderr just before a script exits, and after the context manager has 39 | exited (at least, I think that is why it lets exceptions through). 40 | 41 | ''' 42 | 43 | def __init__(self): 44 | # Open a pair of null files 45 | self.null_fds = [os.open(os.devnull, os.O_RDWR) for x in range(2)] 46 | # Save the actual stdout (1) and stderr (2) file descriptors. 47 | self.save_fds = (os.dup(1), os.dup(2)) 48 | 49 | def __enter__(self): 50 | # Assign the null pointers to stdout and stderr. 51 | os.dup2(self.null_fds[0], 1) 52 | os.dup2(self.null_fds[1], 2) 53 | 54 | def __exit__(self, *_): 55 | # Re-assign the real stdout/stderr back to (1) and (2) 56 | os.dup2(self.save_fds[0], 1) 57 | os.dup2(self.save_fds[1], 2) 58 | # Close the null files 59 | os.close(self.null_fds[0]) 60 | os.close(self.null_fds[1]) 61 | -------------------------------------------------------------------------------- /utils/plot/CommonFigures.py: -------------------------------------------------------------------------------- 1 | from torch import Tensor 2 | import numpy as np 3 | from matplotlib import pyplot as plt 4 | 5 | 6 | def matrix_heatmap(harvest: np.array, save_path='./heatmap_of_matrix.png'): 7 | plt.imshow(harvest) 8 | plt.tight_layout() 9 | plt.colorbar() 10 | plt.axis('off') 11 | plt.savefig(save_path, bbox_inches='tight') 12 | plt.close() 13 | 14 | 15 | def tensor_heatmap(x: Tensor, save_path='./heatmap.png', show=False): 16 | plt.imshow(x.cpu().numpy()) 17 | plt.tight_layout() 18 | plt.colorbar() 19 | plt.axis('off') 20 | if show: 21 | plt.show() 22 | plt.savefig(save_path, bbox_inches='tight') 23 | plt.close() 24 | -------------------------------------------------------------------------------- /utils/plot/Landscape4Model.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.nn as nn 4 | import numpy as np 5 | from tqdm import tqdm 6 | import copy 7 | from matplotlib import pyplot as plt 8 | from mpl_toolkits.mplot3d import Axes3D 9 | from .ColorUtils import get_datetime_str 10 | 11 | modes = ['3D', 'contour', '2D'] 12 | 13 | 14 | class Landscape4Model(): 15 | def __init__(self, model: nn.Module, loss: callable, mode='2D', 16 | save_path='./landscape/'): 17 | ''' 18 | 19 | :param model: 20 | :param loss: given a model, return loss 21 | ''' 22 | self.save_path = save_path 23 | assert mode in modes 24 | self.mode = mode 25 | self.model = model 26 | self.loss = loss 27 | self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 28 | 29 | @torch.no_grad() 30 | def synthesize_coordinates(self, x_min=-0.02, x_max=0.02, x_interval=1e-3, 31 | y_min=-0.02, y_max=0.02, y_interval=1e-3): 32 | x = np.arange(x_min, x_max, x_interval) 33 | y = np.arange(y_min, y_max, y_interval) 34 | self.mesh_x, self.mesh_y = np.meshgrid(x, y) 35 | 36 | @torch.no_grad() 37 | def draw(self): 38 | self._find_direction() 39 | z = self._compute_for_draw() 40 | self.draw_figure(self.mesh_x, self.mesh_y, z) 41 | 42 | @torch.no_grad() 43 | def _find_direction(self): 44 | self.x_unit_vector = {} 45 | self.y_unit_vector = {} 46 | for name, param in self.model.named_parameters(): 47 | self.x_unit_vector[name] = torch.randn_like(param.data) 48 | self.y_unit_vector[name] = torch.randn_like(param.data) 49 | 50 | @torch.no_grad() 51 | def _compute_for_draw(self): 52 | result = [] 53 | if self.mode == '2D': 54 | self.mesh_x = self.mesh_x[0] 55 | for i in tqdm(range(self.mesh_x.shape[0])): 56 | now_x = self.mesh_x[i] 57 | loss = self._compute_loss_for_one_coordinate(now_x, 0) 58 | result.append(loss) 59 | else: 60 | for i in tqdm(range(self.mesh_x.shape[0])): 61 | for j in range(self.mesh_x.shape[1]): 62 | now_x = self.mesh_x[i, j] 63 | now_y = self.mesh_y[i, j] 64 | loss = self._compute_loss_for_one_coordinate(now_x, now_y) 65 | result.append(loss) 66 | 67 | result = np.array(result) 68 | result = result.reshape(self.mesh_x.shape) 69 | return result 70 | 71 | @torch.no_grad() 72 | def _compute_loss_for_one_coordinate(self, now_x: float, now_y: float): 73 | temp_model = copy.deepcopy(self.model) 74 | for name, param in temp_model.named_parameters(): 75 | param.data = param.data + now_x * self.x_unit_vector[name] + now_y * self.y_unit_vector[name] 76 | 77 | return self.loss(temp_model) 78 | 79 | def draw_figure(self, mesh_x, mesh_y, mesh_z): 80 | if self.mode == '3D': 81 | figure = plt.figure() 82 | axes = Axes3D(figure) 83 | axes.plot_surface(mesh_x, mesh_y, mesh_z, cmap='rainbow') 84 | elif self.mode == '2D': 85 | plt.plot(mesh_x, mesh_z) 86 | plt.savefig(os.path.join(self.save_path, get_datetime_str() + ".png")) 87 | plt.close() 88 | 89 | 90 | if __name__ == '__main__': 91 | pass 92 | -------------------------------------------------------------------------------- /utils/plot/__init__.py: -------------------------------------------------------------------------------- 1 | from .Landscape4Input import Landscape4Input 2 | from .Landscape4Model import Landscape4Model 3 | -------------------------------------------------------------------------------- /utils/saver.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | 4 | def print_to_file(x: Any, file_name: str = "log.txt"): 5 | with open(file_name, "w") as f: 6 | print(x, file=f) 7 | -------------------------------------------------------------------------------- /utils/seed.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import random 3 | import numpy as np 4 | 5 | 6 | def set_seed(seed=1): 7 | torch.manual_seed(seed) 8 | random.seed(seed) 9 | np.random.seed(seed) 10 | if torch.cuda.is_available(): 11 | torch.cuda.manual_seed_all(seed) 12 | -------------------------------------------------------------------------------- /utils/types.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | 4 | class EasyDict(dict): 5 | """Convenience class that behaves like a dict but allows access with the attribute syntax.""" 6 | 7 | def __getattr__(self, name: str) -> Any: 8 | try: 9 | return self[name] 10 | except KeyError: 11 | raise AttributeError(name) 12 | 13 | def __setattr__(self, name: str, value: Any) -> None: 14 | self[name] = value 15 | 16 | def __delattr__(self, name: str) -> None: 17 | del self[name] 18 | --------------------------------------------------------------------------------