├── .gitignore ├── requirements.txt ├── data └── indices │ ├── cifar10 │ ├── idx-20.pkl │ ├── idx-40.pkl │ ├── idx-60.pkl │ └── idx-80.pkl │ └── cifar100 │ ├── idx-20.pkl │ ├── idx-40.pkl │ ├── idx-60.pkl │ └── idx-80.pkl ├── models ├── __init__.py ├── vgg.py ├── resnet.py └── densenet.py ├── attacks ├── __init__.py ├── diff_aug_pgd.py ├── pgd_attacker.py └── robust_workers.py ├── utils ├── __init__.py ├── argument.py ├── data.py ├── imagenet_utils.py └── generic.py ├── scripts ├── cifar10 │ ├── gen_em.sh │ ├── gen_tap.sh │ ├── train_clr.sh │ ├── train_rem.sh │ ├── gen_rem.sh │ └── train_rem_mix.sh └── in-mini │ ├── gen_tap.sh │ ├── train_clr.sh │ ├── gen_em.sh │ ├── train_rem.sh │ └── gen_rem.sh ├── README.md ├── LICENSE ├── generate_em.py ├── generate_tap.py ├── train.py └── generate_robust_em.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *~ 3 | *.swp 4 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch==1.8.1 2 | torchvision==0.9.1 3 | opencv-python==4.5.5.62 4 | -------------------------------------------------------------------------------- /data/indices/cifar10/idx-20.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar10/idx-20.pkl -------------------------------------------------------------------------------- /data/indices/cifar10/idx-40.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar10/idx-40.pkl -------------------------------------------------------------------------------- /data/indices/cifar10/idx-60.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar10/idx-60.pkl -------------------------------------------------------------------------------- /data/indices/cifar10/idx-80.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar10/idx-80.pkl -------------------------------------------------------------------------------- /data/indices/cifar100/idx-20.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar100/idx-20.pkl -------------------------------------------------------------------------------- /data/indices/cifar100/idx-40.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar100/idx-40.pkl -------------------------------------------------------------------------------- /data/indices/cifar100/idx-60.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar100/idx-60.pkl -------------------------------------------------------------------------------- /data/indices/cifar100/idx-80.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CausalLearning/robust-unlearnable-examples/HEAD/data/indices/cifar100/idx-80.pkl -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import resnet18, resnet50 2 | from .resnet import wrn34_10 3 | from .vgg import vgg11_bn, vgg16_bn, vgg19_bn 4 | from .densenet import densenet121 5 | -------------------------------------------------------------------------------- /attacks/__init__.py: -------------------------------------------------------------------------------- 1 | from .pgd_attacker import PGDAttacker 2 | from .robust_workers import RobustPGDAttacker, RobustMinimaxPGDDefender 3 | 4 | from .diff_aug_pgd import DiffAugPGDAttacker 5 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .argument import add_shared_args 2 | from .generic import AverageMeter, add_log, generic_init 3 | from .generic import get_indexed_loader, get_indexed_tensor_loader, get_poisoned_loader, get_clear_loader 4 | from .generic import get_arch, get_optim, evaluate 5 | from .generic import get_dataset 6 | from .generic import get_transforms 7 | from .generic import get_model_state 8 | from .data import Dataset, IndexedDataset, IndexedTensorDataset, Loader 9 | -------------------------------------------------------------------------------- /scripts/cifar10/gen_em.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | 3 | python generate_em.py \ 4 | --arch resnet18 \ 5 | --dataset cifar10 \ 6 | --train-steps 5000 \ 7 | --batch-size 128 \ 8 | --optim sgd \ 9 | --lr 0.1 \ 10 | --lr-decay-rate 0.1 \ 11 | --lr-decay-freq 2000 \ 12 | --weight-decay 5e-4 \ 13 | --momentum 0.9 \ 14 | --pgd-radius 8 \ 15 | --pgd-steps 10 \ 16 | --pgd-step-size 1.6 \ 17 | --pgd-random-start \ 18 | --report-freq 1000 \ 19 | --save-freq 1000 \ 20 | --data-dir ./data \ 21 | --save-dir ./exp_data/cifar10/noise/em8 \ 22 | --save-name em 23 | -------------------------------------------------------------------------------- /scripts/cifar10/gen_tap.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | 3 | python generate_tap.py \ 4 | --arch resnet18 \ 5 | --dataset cifar10 \ 6 | --targeted \ 7 | --batch-size 128 \ 8 | --optim sgd \ 9 | --lr 0.1 \ 10 | --lr-decay-rate 0.1 \ 11 | --lr-decay-freq 2000 \ 12 | --pgd-radius 8 \ 13 | --pgd-steps 250 \ 14 | --pgd-step-size 0.064 \ 15 | --pgd-random-start \ 16 | --samp-num 1 \ 17 | --resume \ 18 | --resume-path ./exp_data/cifar10/train/clr/r0/train-fin-model.pkl \ 19 | --data-dir ./data \ 20 | --save-dir ./exp_data/cifar10/noise/tap8 \ 21 | --save-name tap 22 | -------------------------------------------------------------------------------- /scripts/cifar10/train_clr.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | 3 | python train.py \ 4 | --arch resnet18 \ 5 | --dataset cifar10 \ 6 | --train-steps 40000 \ 7 | --batch-size 128 \ 8 | --optim sgd \ 9 | --lr 0.1 \ 10 | --lr-decay-rate 0.1 \ 11 | --lr-decay-freq 16000 \ 12 | --weight-decay 5e-4 \ 13 | --momentum 0.9 \ 14 | --pgd-radius 0 \ 15 | --pgd-steps 10 \ 16 | --pgd-step-size 0 \ 17 | --pgd-random-start \ 18 | --report-freq 1000 \ 19 | --save-freq 100000 \ 20 | --data-dir ./data \ 21 | --save-dir ./exp_data/cifar10/train/clr/r0 \ 22 | --save-name train 23 | -------------------------------------------------------------------------------- /scripts/in-mini/gen_tap.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | imagenet_path=$2 3 | 4 | python generate_tap.py \ 5 | --arch resnet18 \ 6 | --dataset imagenet-mini \ 7 | --targeted \ 8 | --batch-size 128 \ 9 | --optim sgd \ 10 | --lr 0.1 \ 11 | --lr-decay-rate 0.1 \ 12 | --lr-decay-freq 2000 \ 13 | --pgd-radius 8 \ 14 | --pgd-steps 100 \ 15 | --pgd-step-size 0.16 \ 16 | --pgd-random-start \ 17 | --samp-num 1 \ 18 | --parallel \ 19 | --resume \ 20 | --resume-path ./exp_data/in-mini/train/clr/r0/train-fin-model.pkl \ 21 | --data-dir ${imagenet_path} \ 22 | --save-dir ./exp_data/in-mini/noise/tap8 \ 23 | --save-name tap 24 | -------------------------------------------------------------------------------- /scripts/in-mini/train_clr.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | imagenet_path=$2 3 | 4 | python train.py \ 5 | --arch resnet18 \ 6 | --dataset imagenet-mini \ 7 | --train-steps 40000 \ 8 | --batch-size 128 \ 9 | --optim sgd \ 10 | --lr 0.1 \ 11 | --lr-decay-rate 0.1 \ 12 | --lr-decay-freq 16000 \ 13 | --weight-decay 5e-4 \ 14 | --momentum 0.9 \ 15 | --pgd-radius 0 \ 16 | --pgd-steps 8 \ 17 | --pgd-step-size 0 \ 18 | --pgd-random-start \ 19 | --report-freq 1000 \ 20 | --save-freq 100000 \ 21 | --parallel \ 22 | --data-dir ${imagenet_path} \ 23 | --save-dir ./exp_data/in-mini/train/clr/r0 \ 24 | --save-name train 25 | -------------------------------------------------------------------------------- /scripts/in-mini/gen_em.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | imagenet_path=$2 3 | 4 | python generate_em.py \ 5 | --arch resnet18 \ 6 | --dataset imagenet-mini \ 7 | --train-steps 3000 \ 8 | --batch-size 128 \ 9 | --optim sgd \ 10 | --lr 0.1 \ 11 | --lr-decay-rate 0.1 \ 12 | --lr-decay-freq 1200 \ 13 | --weight-decay 5e-4 \ 14 | --momentum 0.9 \ 15 | --pgd-radius 8 \ 16 | --pgd-steps 7 \ 17 | --pgd-step-size 2 \ 18 | --pgd-random-start \ 19 | --perturb-freq 1 \ 20 | --report-freq 500 \ 21 | --save-freq 500 \ 22 | --parallel \ 23 | --data-dir ${imagenet_path} \ 24 | --save-dir ./exp_data/in-mini/noise/em8 \ 25 | --save-name em 26 | -------------------------------------------------------------------------------- /scripts/cifar10/train_rem.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | 3 | python train.py \ 4 | --arch resnet18 \ 5 | --dataset cifar10 \ 6 | --train-steps 40000 \ 7 | --batch-size 128 \ 8 | --optim sgd \ 9 | --lr 0.1 \ 10 | --lr-decay-rate 0.1 \ 11 | --lr-decay-freq 16000 \ 12 | --weight-decay 5e-4 \ 13 | --momentum 0.9 \ 14 | --pgd-radius 4 \ 15 | --pgd-steps 10 \ 16 | --pgd-step-size 0.8 \ 17 | --pgd-random-start \ 18 | --report-freq 1000 \ 19 | --save-freq 100000 \ 20 | --noise-path ./exp_data/cifar10/noise/rem8-4/rem-fin-def-noise.pkl \ 21 | --data-dir ./data \ 22 | --save-dir ./exp_data/cifar10/train/rem8-4/r4 \ 23 | --save-name train 24 | -------------------------------------------------------------------------------- /scripts/in-mini/train_rem.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | imagenet_path=$2 3 | 4 | python train.py \ 5 | --arch resnet18 \ 6 | --dataset imagenet-mini \ 7 | --train-steps 40000 \ 8 | --batch-size 128 \ 9 | --optim sgd \ 10 | --lr 0.1 \ 11 | --lr-decay-rate 0.1 \ 12 | --lr-decay-freq 16000 \ 13 | --weight-decay 5e-4 \ 14 | --momentum 0.9 \ 15 | --pgd-radius 4 \ 16 | --pgd-steps 8 \ 17 | --pgd-step-size 1 \ 18 | --pgd-random-start \ 19 | --report-freq 1000 \ 20 | --save-freq 100000 \ 21 | --parallel \ 22 | --data-dir ${imagenet_path} \ 23 | --noise-path ./exp_data/in-mini/noise/rem8-4/rem-fin-def-noise.pkl \ 24 | --save-dir ./exp_data/in-mini/train/rem8-4/r4 \ 25 | --save-name train 26 | -------------------------------------------------------------------------------- /scripts/cifar10/gen_rem.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | 3 | python generate_robust_em.py \ 4 | --arch resnet18 \ 5 | --dataset cifar10 \ 6 | --train-steps 5000 \ 7 | --batch-size 128 \ 8 | --optim sgd \ 9 | --lr 0.1 \ 10 | --lr-decay-rate 0.1 \ 11 | --lr-decay-freq 2000 \ 12 | --weight-decay 5e-4 \ 13 | --momentum 0.9 \ 14 | --pgd-radius 8 \ 15 | --pgd-steps 10 \ 16 | --pgd-step-size 1.6 \ 17 | --pgd-random-start \ 18 | --atk-pgd-radius 4 \ 19 | --atk-pgd-steps 10 \ 20 | --atk-pgd-step-size 0.8 \ 21 | --atk-pgd-random-start \ 22 | --samp-num 5 \ 23 | --report-freq 1000 \ 24 | --save-freq 1000 \ 25 | --data-dir ./data \ 26 | --save-dir ./exp_data/cifar10/noise/rem8-4 \ 27 | --save-name rem 28 | -------------------------------------------------------------------------------- /scripts/cifar10/train_rem_mix.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | 3 | for i in {2..8..2} 4 | do 5 | python train.py \ 6 | --data-mode mix \ 7 | --arch resnet18 \ 8 | --dataset cifar10 \ 9 | --train-steps 40000 \ 10 | --batch-size 128 \ 11 | --optim sgd \ 12 | --lr 0.1 \ 13 | --lr-decay-rate 0.1 \ 14 | --lr-decay-freq 16000 \ 15 | --weight-decay 5e-4 \ 16 | --momentum 0.9 \ 17 | --pgd-radius 4 \ 18 | --pgd-steps 10 \ 19 | --pgd-step-size 0.8 \ 20 | --pgd-random-start \ 21 | --report-freq 1000 \ 22 | --save-freq 100000 \ 23 | --poi-idx-path ./data/indices/cifar10/idx-${i}0.pkl \ 24 | --noise-path ./exp_data/cifar10/noise/rem8-4/rem-fin-def-noise.pkl \ 25 | --data-dir ./data \ 26 | --save-dir ./exp_data/cifar10/mixing/rem8-4/${i}0/r4 \ 27 | --save-name train 28 | done 29 | -------------------------------------------------------------------------------- /scripts/in-mini/gen_rem.sh: -------------------------------------------------------------------------------- 1 | cd $1 2 | imagenet_path=$2 3 | 4 | python generate_robust_em.py \ 5 | --arch resnet18 \ 6 | --dataset imagenet-mini \ 7 | --train-steps 3000 \ 8 | --batch-size 128 \ 9 | --optim sgd \ 10 | --lr 0.1 \ 11 | --lr-decay-rate 0.1 \ 12 | --lr-decay-freq 1200 \ 13 | --weight-decay 5e-4 \ 14 | --momentum 0.9 \ 15 | --pgd-radius 8 \ 16 | --pgd-steps 7 \ 17 | --pgd-step-size 2 \ 18 | --pgd-random-start \ 19 | --atk-pgd-radius 4 \ 20 | --atk-pgd-steps 7 \ 21 | --atk-pgd-step-size 1 \ 22 | --atk-pgd-random-start \ 23 | --samp-num 4 \ 24 | --perturb-freq 1 \ 25 | --report-freq 500 \ 26 | --save-freq 500 \ 27 | --parallel \ 28 | --data-dir ${imagenet_path} \ 29 | --save-dir ./exp_data/in-mini/noise/rem8-4 \ 30 | --save-name rem 31 | -------------------------------------------------------------------------------- /attacks/diff_aug_pgd.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class DiffAugPGDAttacker(): 5 | def __init__(self, samp_num, trans, 6 | radius, steps, step_size, random_start, ascending=True): 7 | self.samp_num = samp_num 8 | self.trans = trans 9 | 10 | self.radius = radius / 255. 11 | self.steps = steps 12 | self.step_size = step_size / 255. 13 | self.random_start = random_start 14 | self.ascending = ascending 15 | 16 | def perturb(self, model, criterion, x, y): 17 | ''' initialize noise ''' 18 | delta = torch.zeros_like(x.data) 19 | if self.steps==0 or self.radius==0: 20 | return delta 21 | 22 | if self.random_start: 23 | delta.uniform_(-self.radius, self.radius) 24 | 25 | ''' temporarily shutdown autograd of model to improve pgd efficiency ''' 26 | model.eval() 27 | for pp in model.parameters(): 28 | pp.requires_grad = False 29 | 30 | delta.requires_grad_() 31 | for step in range(self.steps): 32 | delta.grad = None 33 | 34 | for i in range(self.samp_num): 35 | adv_x = (self.trans(x) + delta).clamp(0., 255.) 36 | _y = model(adv_x) 37 | lo = criterion(_y, y) 38 | lo.backward() 39 | 40 | with torch.no_grad(): 41 | grad = delta.grad.data 42 | if not self.ascending: grad.mul_(-1) 43 | delta.add_(torch.sign(grad), alpha=self.step_size) 44 | delta.clamp_(-self.radius, self.radius) 45 | 46 | ''' reopen autograd of model after pgd ''' 47 | for pp in model.parameters(): 48 | pp.requires_grad = True 49 | 50 | return delta.data 51 | -------------------------------------------------------------------------------- /attacks/pgd_attacker.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class PGDAttacker(): 5 | def __init__(self, radius, steps, step_size, random_start, norm_type, ascending=True): 6 | self.radius = radius / 255. 7 | self.steps = steps 8 | self.step_size = step_size / 255. 9 | self.random_start = random_start 10 | self.norm_type = norm_type 11 | self.ascending = ascending 12 | 13 | def perturb(self, model, criterion, x, y): 14 | if self.steps==0 or self.radius==0: 15 | return x.clone() 16 | 17 | adv_x = x.clone() 18 | if self.random_start: 19 | if self.norm_type == 'l-infty': 20 | adv_x += 2 * (torch.rand_like(x) - 0.5) * self.radius 21 | else: 22 | adv_x += 2 * (torch.rand_like(x) - 0.5) * self.radius / self.steps 23 | self._clip_(adv_x, x) 24 | 25 | ''' temporarily shutdown autograd of model to improve pgd efficiency ''' 26 | model.eval() 27 | for pp in model.parameters(): 28 | pp.requires_grad = False 29 | 30 | for step in range(self.steps): 31 | adv_x.requires_grad_() 32 | _y = model(adv_x) 33 | loss = criterion(_y, y) 34 | grad = torch.autograd.grad(loss, [adv_x])[0] 35 | 36 | with torch.no_grad(): 37 | if not self.ascending: grad.mul_(-1) 38 | 39 | if self.norm_type == 'l-infty': 40 | adv_x.add_(torch.sign(grad), alpha=self.step_size) 41 | else: 42 | if self.norm_type == 'l2': 43 | grad_norm = (grad.reshape(grad.shape[0],-1)**2).sum(dim=1).sqrt() 44 | elif self.norm_type == 'l1': 45 | grad_norm = grad.reshape(grad.shape[0],-1).abs().sum(dim=1) 46 | grad_norm = grad_norm.reshape( -1, *( [1] * (len(x.shape)-1) ) ) 47 | scaled_grad = grad / (grad_norm + 1e-10) 48 | adv_x.add_(scaled_grad, alpha=self.step_size) 49 | 50 | self._clip_(adv_x, x) 51 | 52 | ''' reopen autograd of model after pgd ''' 53 | for pp in model.parameters(): 54 | pp.requires_grad = True 55 | 56 | return adv_x.data 57 | 58 | def _clip_(self, adv_x, x): 59 | adv_x -= x 60 | if self.norm_type == 'l-infty': 61 | adv_x.clamp_(-self.radius, self.radius) 62 | else: 63 | if self.norm_type == 'l2': 64 | norm = (adv_x.reshape(adv_x.shape[0],-1)**2).sum(dim=1).sqrt() 65 | elif self.norm_type == 'l1': 66 | norm = adv_x.reshape(adv_x.shape[0],-1).abs().sum(dim=1) 67 | norm = norm.reshape( -1, *( [1] * (len(x.shape)-1) ) ) 68 | adv_x /= (norm + 1e-10) 69 | adv_x *= norm.clamp(max=self.radius) 70 | adv_x += x 71 | adv_x.clamp_(-0.5, 0.5) 72 | -------------------------------------------------------------------------------- /utils/argument.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | def add_shared_args(parser): 5 | assert isinstance(parser, argparse.ArgumentParser) 6 | 7 | parser.add_argument('--arch', type=str, default='resnet18', 8 | choices=['resnet18', 'resnet50', 'vgg11-bn', 'vgg16-bn', 'vgg19-bn', 'densenet-121', 'inception-resnet-v1', 'inception-v3', 'wrn-34-10'], 9 | help='choose the model architecture') 10 | parser.add_argument('--dataset', type=str, default='cifar10', 11 | choices=['cifar10', 'cifar100', 'tiny-imagenet', 'imagenet-mini'], 12 | help='choose the dataset') 13 | parser.add_argument('--train-steps', type=int, default=80000, 14 | help='set the training steps') 15 | parser.add_argument('--batch-size', type=int, default=128, 16 | help='set the batch size') 17 | 18 | parser.add_argument('--optim', type=str, default='sgd', 19 | choices=['sgd', 'adam'], 20 | help='select which optimizer to use') 21 | parser.add_argument('--lr', type=float, default=0.1, 22 | help='set the initial learning rate') 23 | parser.add_argument('--lr-decay-rate', type=float, default=0.1, 24 | help='set the learning rate decay rate') 25 | parser.add_argument('--lr-decay-freq', type=int, default=30000, 26 | help='set the learning rate decay frequency') 27 | parser.add_argument('--weight-decay', type=float, default=1e-4, 28 | help='set the weight decay rate') 29 | parser.add_argument('--momentum', type=float, default=0.9, 30 | help='set the momentum for SGD') 31 | 32 | parser.add_argument('--pgd-radius', type=float, default=0, 33 | help='set the perturbation radius in pgd') 34 | parser.add_argument('--pgd-steps', type=int, default=0, 35 | help='set the number of iteration steps in pgd') 36 | parser.add_argument('--pgd-step-size', type=float, default=0, 37 | help='set the step size in pgd') 38 | parser.add_argument('--pgd-random-start', action='store_true', 39 | help='if select, randomly choose starting points each time performing pgd') 40 | parser.add_argument('--pgd-norm-type', type=str, default='l-infty', 41 | choices=['l-infty', 'l2', 'l1'], 42 | help='set the type of metric norm in pgd') 43 | 44 | parser.add_argument('--parallel', action='store_true', 45 | help='select to use distributed data parallel') 46 | parser.add_argument('--local_rank', type=int, default=0, 47 | help='for distributed data parallel') 48 | 49 | parser.add_argument('--cpu', action='store_true', 50 | help='select to use cpu, otherwise use gpu') 51 | parser.add_argument('--data-dir', type=str, default='./data', 52 | help='set the path to the exp data') 53 | parser.add_argument('--save-dir', type=str, default='./temp', 54 | help='set which dictionary to save the experiment result') 55 | parser.add_argument('--save-name', type=str, default='temp-name', 56 | help='set the save name of the experiment result') 57 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Robust Unlearnable Examples: Protecting Data Against Adversarial Learning 2 | 3 | This is the official repository for ICLR 2022 paper ["Robust Unlearnable Examples: Protecting Data Against Adversarial Learning"](https://openreview.net/forum?id=baUQQPwQiAg) by Shaopeng Fu, Fengxiang He, Yang Liu, Li Shen and Dacheng Tao. 4 | 5 | ## Requirements 6 | 7 | - Python 3.8 8 | - PyTorch 1.8.1 9 | - Torchvision 0.9.1 10 | - OpenCV 4.5.5 11 | 12 | #### Install dependencies using pip 13 | 14 | ```shell 15 | pip install -r requirements.txt 16 | ``` 17 | 18 | #### Install dependencies using Anaconda 19 | 20 | It is recommended to create your experiment environment with [Anaconda3](https://www.anaconda.com/). 21 | 22 | ```shell 23 | conda install pytorch=1.8.1 torchvision=0.9.1 cudatoolkit=10.2 -c pytorch 24 | conda install -c conda-forge opencv=4.5.5 25 | ``` 26 | 27 | ## Quick Start 28 | 29 | We give an example of creating robust unlearnable examples from CIFAR-10 dataset. More experiment examples can be found in [./scripts](./scripts). 30 | 31 | #### Generate robust error-minimizing noise for CIFAR-10 dataset 32 | 33 | ```bash 34 | python generate_robust_em.py \ 35 | --arch resnet18 \ 36 | --dataset cifar10 \ 37 | --train-steps 5000 \ 38 | --batch-size 128 \ 39 | --optim sgd \ 40 | --lr 0.1 \ 41 | --lr-decay-rate 0.1 \ 42 | --lr-decay-freq 2000 \ 43 | --weight-decay 5e-4 \ 44 | --momentum 0.9 \ 45 | --pgd-radius 8 \ 46 | --pgd-steps 10 \ 47 | --pgd-step-size 1.6 \ 48 | --pgd-random-start \ 49 | --atk-pgd-radius 4 \ 50 | --atk-pgd-steps 10 \ 51 | --atk-pgd-step-size 0.8 \ 52 | --atk-pgd-random-start \ 53 | --samp-num 5 \ 54 | --report-freq 1000 \ 55 | --save-freq 1000 \ 56 | --data-dir ./data \ 57 | --save-dir ./exp_data/cifar10/noise/rem8-4 \ 58 | --save-name rem 59 | ``` 60 | 61 | #### Perform adversarial training on robust unlearnable examples 62 | 63 | ```bash 64 | python train.py \ 65 | --arch resnet18 \ 66 | --dataset cifar10 \ 67 | --train-steps 40000 \ 68 | --batch-size 128 \ 69 | --optim sgd \ 70 | --lr 0.1 \ 71 | --lr-decay-rate 0.1 \ 72 | --lr-decay-freq 16000 \ 73 | --weight-decay 5e-4 \ 74 | --momentum 0.9 \ 75 | --pgd-radius 4 \ 76 | --pgd-steps 10 \ 77 | --pgd-step-size 0.8 \ 78 | --pgd-random-start \ 79 | --report-freq 1000 \ 80 | --save-freq 100000 \ 81 | --noise-path ./exp_data/cifar10/noise/rem8-4/rem-fin-def-noise.pkl \ 82 | --data-dir ./data \ 83 | --save-dir ./exp_data/cifar10/train/rem8-4/r4 \ 84 | --save-name train 85 | ``` 86 | 87 | ## Citation 88 | 89 | ``` 90 | @inproceedings{fu2022robust, 91 | title={Robust Unlearnable Examples: Protecting Data Against Adversarial Learning}, 92 | author={Shaopeng Fu and Fengxiang He and Yang Liu and Li Shen and Dacheng Tao}, 93 | booktitle={International Conference on Learning Representations}, 94 | year={2022} 95 | } 96 | ``` 97 | 98 | ## Acknowledgment 99 | 100 | - Unlearnable examples: [https://github.com/HanxunH/Unlearnable-Examples](https://github.com/HanxunH/Unlearnable-Examples) 101 | - Adversarial poisons: [https://github.com/lhfowl/adversarial_poisons](https://github.com/lhfowl/adversarial_poisons) 102 | - Neural tangent generalization attacks: [https://github.com/lionelmessi6410/ntga](https://github.com/lionelmessi6410/ntga) -------------------------------------------------------------------------------- /models/vgg.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Modified from https://github.com/pytorch/vision.git 3 | ''' 4 | import math 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.init as init 9 | 10 | __all__ = [ 11 | 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 12 | 'vgg19_bn', 'vgg19', 13 | ] 14 | 15 | 16 | class VGG(nn.Module): 17 | ''' VGG model ''' 18 | def __init__(self, features, out_channels): 19 | super(VGG, self).__init__() 20 | self.features = features 21 | self.avgpool = nn.AdaptiveAvgPool2d((1,1)) 22 | self.classifier = nn.Sequential( 23 | nn.Dropout(), 24 | nn.Linear(512, 512), 25 | nn.ReLU(True), 26 | nn.Dropout(), 27 | nn.Linear(512, 512), 28 | nn.ReLU(True), 29 | nn.Linear(512, out_channels), 30 | ) 31 | 32 | ''' Initialize weights ''' 33 | for m in self.modules(): 34 | if isinstance(m, nn.Conv2d): 35 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 36 | m.weight.data.normal_(0, math.sqrt(2. / n)) 37 | m.bias.data.zero_() 38 | 39 | 40 | def forward(self, x): 41 | x = self.features(x) 42 | x = self.avgpool(x) 43 | x = torch.flatten(x, 1) 44 | x = self.classifier(x) 45 | return x 46 | 47 | 48 | def make_layers(cfg, in_dims=3, batch_norm=False): 49 | layers = [] 50 | in_channels = in_dims 51 | for v in cfg: 52 | if v == 'M': 53 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 54 | else: 55 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 56 | if batch_norm: 57 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 58 | else: 59 | layers += [conv2d, nn.ReLU(inplace=True)] 60 | in_channels = v 61 | return nn.Sequential(*layers) 62 | 63 | 64 | cfg = { 65 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 66 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 67 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 68 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 69 | } 70 | 71 | 72 | # def vgg11(out_channels=10): 73 | # """VGG 11-layer model (configuration "A")""" 74 | # return VGG(make_layers(cfg['A']), out_channels) 75 | # 76 | # 77 | def vgg11_bn(in_dims=3, out_dims=10): 78 | """VGG 11-layer model (configuration "A") with batch normalization""" 79 | return VGG(make_layers(cfg['A'], in_dims, batch_norm=True), out_dims) 80 | # 81 | # 82 | # def vgg13(out_channels=10): 83 | # """VGG 13-layer model (configuration "B")""" 84 | # return VGG(make_layers(cfg['B']), out_channels) 85 | # 86 | # 87 | # def vgg13_bn(out_channels=10): 88 | # """VGG 13-layer model (configuration "B") with batch normalization""" 89 | # return VGG(make_layers(cfg['B'], batch_norm=True), out_channels) 90 | 91 | 92 | # def vgg16(out_dims=10): 93 | # """VGG 16-layer model (configuration "D")""" 94 | # return VGG(make_layers(cfg['D']), out_dims) 95 | 96 | 97 | def vgg16_bn(in_dims=3, out_dims=10): 98 | """VGG 16-layer model (configuration "D") with batch normalization""" 99 | return VGG(make_layers(cfg['D'], in_dims, batch_norm=True), out_dims) 100 | 101 | 102 | # def vgg19(out_channels=10): 103 | # """VGG 19-layer model (configuration "E")""" 104 | # return VGG(make_layers(cfg['E']), out_channels) 105 | 106 | 107 | def vgg19_bn(in_dims=3, out_dims=10): 108 | """VGG 19-layer model (configuration 'E') with batch normalization""" 109 | return VGG(make_layers(cfg['E'], in_dims, batch_norm=True), out_dims) 110 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Jingdong Technology Information Technology Co., Ltd. 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. 22 | 23 | ====== 24 | 25 | MIT License 26 | 27 | Copyright (c) 2017 liukuang 28 | 29 | Permission is hereby granted, free of charge, to any person obtaining a copy 30 | of this software and associated documentation files (the "Software"), to deal 31 | in the Software without restriction, including without limitation the rights 32 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 33 | copies of the Software, and to permit persons to whom the Software is 34 | furnished to do so, subject to the following conditions: 35 | 36 | The above copyright notice and this permission notice shall be included in all 37 | copies or substantial portions of the Software. 38 | 39 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 40 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 41 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 42 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 43 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 44 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 45 | SOFTWARE. 46 | 47 | ====== 48 | 49 | BSD 3-Clause License 50 | 51 | Copyright (c) Soumith Chintala 2016, 52 | All rights reserved. 53 | 54 | Redistribution and use in source and binary forms, with or without 55 | modification, are permitted provided that the following conditions are met: 56 | 57 | * Redistributions of source code must retain the above copyright notice, this 58 | list of conditions and the following disclaimer. 59 | 60 | * Redistributions in binary form must reproduce the above copyright notice, 61 | this list of conditions and the following disclaimer in the documentation 62 | and/or other materials provided with the distribution. 63 | 64 | * Neither the name of the copyright holder nor the names of its 65 | contributors may be used to endorse or promote products derived from 66 | this software without specific prior written permission. 67 | 68 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 69 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 70 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 71 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 72 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 73 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 74 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 75 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 76 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 77 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 78 | -------------------------------------------------------------------------------- /utils/data.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import numpy as np 3 | import torch 4 | import torchvision 5 | import os 6 | import pickle 7 | 8 | 9 | class ElementWiseTransform(): 10 | def __init__(self, trans=None): 11 | self.trans = trans 12 | 13 | def __call__(self, x): 14 | if self.trans is None: return x 15 | return torch.cat( [self.trans( xx.view(1, *xx.shape) ) for xx in x] ) 16 | 17 | 18 | class IndexedTensorDataset(): 19 | def __init__(self, x, y): 20 | self.x = x 21 | self.y = y 22 | 23 | def __getitem__(self, idx): 24 | x, y = self.x[idx], self.y[idx] 25 | ''' transform HWC pic to CWH pic ''' 26 | x = torch.tensor(x, dtype=torch.float32).permute(2,0,1) 27 | return x, y, idx 28 | 29 | def __len__(self): 30 | return len(self.x) 31 | 32 | 33 | class Dataset(): 34 | def __init__(self, x, y, transform=None, fitr=None): 35 | self.x = x 36 | self.y = y 37 | self.transform = transform 38 | self.fitr = fitr 39 | 40 | def __getitem__(self, idx): 41 | x, y = self.x[idx], self.y[idx] 42 | 43 | ''' low pass filtering ''' 44 | if self.fitr is not None: 45 | x = self.fitr(x) 46 | 47 | ''' data augmentation ''' 48 | if self.transform is not None: 49 | x = self.transform( Image.fromarray(x) ) 50 | 51 | return x, y 52 | 53 | def __len__(self): 54 | return len(self.x) 55 | 56 | 57 | class IndexedDataset(): 58 | def __init__(self, x, y, transform=None): 59 | self.x = x 60 | self.y = y 61 | self.ii = np.array( range(len(x)), dtype=np.int64 ) 62 | self.transform = transform 63 | 64 | def __getitem__(self, idx): 65 | x, y, ii = Image.fromarray(self.x[idx]), self.y[idx], self.ii[idx] 66 | if self.transform is not None: 67 | x = self.transform(x) 68 | return x, y, ii 69 | 70 | def __len__(self): 71 | return len(self.x) 72 | 73 | 74 | def datasetCIFAR10(root='./path', train=True, transform=None): 75 | return torchvision.datasets.CIFAR10(root=root, train=train, 76 | transform=transform, download=True) 77 | 78 | def datasetCIFAR100(root='./path', train=True, transform=None): 79 | return torchvision.datasets.CIFAR100(root=root, train=train, 80 | transform=transform, download=True) 81 | 82 | def datasetTinyImageNet(root='./path', train=True, transform=None): 83 | if train: root = os.path.join(root, 'tiny-imagenet_train.pkl') 84 | else: root = os.path.join(root, 'tiny-imagenet_val.pkl') 85 | with open(root, 'rb') as f: 86 | dat = pickle.load(f) 87 | return Dataset(dat['data'], dat['targets'], transform) 88 | # root = os.path.join(root, 'tiny-imagenet-200') 89 | # if train: root = os.path.join(root, 'train') 90 | # else: root = os.path.join(root, 'val', 'images') 91 | # raw_dataset = torchvision.datasets.ImageFolder(root) 92 | # xx, yy = [], [] 93 | # for i in range( len(raw_dataset) ): 94 | # x, y = raw_dataset[i] 95 | # x = np.array(x) 96 | # xx.append( x.reshape(1, *x.shape) ) 97 | # yy.append( y ) 98 | # xx = np.concatenate(xx) 99 | # yy = np.array(yy) 100 | 101 | # dat = {'data':xx, 'targets':yy} 102 | # if train: save_name = 'tiny-imagenet_train.pkl' 103 | # else: save_name = 'tiny-imagenet_val.pkl' 104 | 105 | # import pickle 106 | # with open('./data/{}'.format(save_name), 'wb') as f: 107 | # pickle.dump(dat, f) 108 | # exit() 109 | # return Dataset(xx, yy, transform) 110 | 111 | 112 | class Loader(): 113 | def __init__(self, dataset, batch_size, shuffle=False, drop_last=False, num_workers=4): 114 | self.loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers) 115 | self.iterator = None 116 | 117 | def __iter__(self): 118 | return iter(self.loader) 119 | 120 | def __len__(self): 121 | return len(self.loader) 122 | 123 | def __next__(self): 124 | if self.iterator is None: 125 | self.iterator = iter(self.loader) 126 | 127 | try: 128 | samples = next(self.iterator) 129 | except StopIteration: 130 | self.iterator = iter(self.loader) 131 | samples = next(self.iterator) 132 | 133 | return samples 134 | -------------------------------------------------------------------------------- /attacks/robust_workers.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class RobustPGDAttacker(): 5 | def __init__(self, samp_num, trans, 6 | radius, steps, step_size, random_start, ascending=True): 7 | self.samp_num = samp_num 8 | self.trans = trans 9 | 10 | self.radius = radius / 255. 11 | self.steps = steps 12 | self.step_size = step_size / 255. 13 | self.random_start = random_start 14 | self.ascending = ascending 15 | 16 | def perturb(self, model, criterion, x, y): 17 | ''' initialize noise ''' 18 | delta = torch.zeros_like(x.data) 19 | if self.steps==0 or self.radius==0: 20 | return delta 21 | 22 | if self.random_start: 23 | delta.uniform_(-self.radius, self.radius) 24 | 25 | ''' temporarily shutdown autograd of model to improve pgd efficiency ''' 26 | model.eval() 27 | for pp in model.parameters(): 28 | pp.requires_grad = False 29 | 30 | delta.requires_grad_() 31 | for step in range(self.steps): 32 | delta.grad = None 33 | 34 | for i in range(self.samp_num): 35 | adv_x = self.trans( (x + delta * 255).clamp(0., 255.) ) 36 | _y = model(adv_x) 37 | lo = criterion(_y, y) 38 | lo.backward() 39 | 40 | with torch.no_grad(): 41 | grad = delta.grad.data 42 | if not self.ascending: grad.mul_(-1) 43 | delta.add_(torch.sign(grad), alpha=self.step_size) 44 | delta.clamp_(-self.radius, self.radius) 45 | 46 | ''' reopen autograd of model after pgd ''' 47 | for pp in model.parameters(): 48 | pp.requires_grad = True 49 | 50 | return delta.data 51 | 52 | 53 | class RobustMinimaxPGDDefender(): 54 | def __init__(self, samp_num, trans, 55 | radius, steps, step_size, random_start, 56 | atk_radius, atk_steps, atk_step_size, atk_random_start): 57 | self.samp_num = samp_num 58 | self.trans = trans 59 | 60 | self.radius = radius / 255. 61 | self.steps = steps 62 | self.step_size = step_size / 255. 63 | self.random_start = random_start 64 | 65 | self.atk_radius = atk_radius / 255. 66 | self.atk_steps = atk_steps 67 | self.atk_step_size = atk_step_size / 255. 68 | self.atk_random_start = atk_random_start 69 | 70 | def perturb(self, model, criterion, x, y): 71 | ''' initialize noise ''' 72 | delta = torch.zeros_like(x) 73 | if self.steps==0 or self.radius==0: 74 | return delta 75 | 76 | if self.random_start: 77 | delta.uniform_(-self.radius, self.radius) 78 | 79 | ''' temporarily disable autograd of model to improve pgd efficiency ''' 80 | model.eval() 81 | for pp in model.parameters(): 82 | pp.requires_grad = False 83 | 84 | delta.requires_grad_() 85 | for step in range(self.steps): 86 | delta.grad = None 87 | 88 | for i in range(self.samp_num): 89 | def_x = self.trans( (x + delta * 255).clamp(0., 255.) ) 90 | adv_x = self._get_adv_(model, criterion, def_x.data, y) 91 | 92 | adv_x.requires_grad_() 93 | _y = model(adv_x) 94 | lo = criterion(_y, y) 95 | 96 | gd = torch.autograd.grad(lo, [adv_x])[0] 97 | 98 | upd_lo = (def_x * gd).sum() 99 | upd_lo.backward() 100 | 101 | with torch.no_grad(): 102 | grad = delta.grad.data 103 | grad.mul_(-1) 104 | delta.add_(torch.sign(grad), alpha=self.step_size) 105 | delta.clamp_(-self.radius, self.radius) 106 | 107 | ''' re-enable autograd of model after pgd ''' 108 | for pp in model.parameters(): 109 | pp.requires_grad = True 110 | 111 | return delta.data 112 | 113 | def _get_adv_(self, model, criterion, x, y): 114 | adv_x = x.clone() 115 | if self.atk_steps==0 or self.atk_radius==0: 116 | return adv_x 117 | 118 | if self.atk_random_start: 119 | adv_x += 2 * (torch.rand_like(x) - 0.5) * self.atk_radius 120 | self._clip_(adv_x, x, radius=self.atk_radius) 121 | 122 | for step in range(self.atk_steps): 123 | adv_x.requires_grad_() 124 | _y = model(adv_x) 125 | loss = criterion(_y, y) 126 | 127 | ''' gradient ascent ''' 128 | grad = torch.autograd.grad(loss, [adv_x])[0] 129 | 130 | with torch.no_grad(): 131 | adv_x.add_(torch.sign(grad), alpha=self.atk_step_size) 132 | self._clip_(adv_x, x, radius=self.atk_radius) 133 | 134 | return adv_x.data 135 | 136 | def _clip_(self, adv_x, x, radius): 137 | adv_x -= x 138 | adv_x.clamp_(-radius, radius) 139 | adv_x += x 140 | adv_x.clamp_(-0.5, 0.5) 141 | -------------------------------------------------------------------------------- /generate_em.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import argparse 4 | import numpy as np 5 | import torch 6 | 7 | import utils 8 | import attacks 9 | 10 | 11 | def get_args(): 12 | parser = argparse.ArgumentParser() 13 | utils.add_shared_args(parser) 14 | 15 | parser.add_argument('--perturb-freq', type=int, default=1, 16 | help='set the perturbation frequency') 17 | parser.add_argument('--report-freq', type=int, default=500, 18 | help='set the report frequency') 19 | parser.add_argument('--save-freq', type=int, default=5000, 20 | help='set the checkpoint saving frequency') 21 | 22 | return parser.parse_args() 23 | 24 | 25 | def regenerate_def_noise(def_noise, model, criterion, loader, defender, cpu): 26 | for x, y, ii in loader: 27 | if not cpu: x, y = x.cuda(), y.cuda() 28 | def_x = defender.perturb(model, criterion, x, y) 29 | def_noise[ii] = (def_x - x).cpu().numpy() 30 | 31 | 32 | def save_checkpoint(save_dir, save_name, model, optim, log, def_noise=None): 33 | torch.save({ 34 | 'model_state_dict': utils.get_model_state(model), 35 | 'optim_state_dict': optim.state_dict(), 36 | }, os.path.join(save_dir, '{}-model.pkl'.format(save_name))) 37 | with open(os.path.join(save_dir, '{}-log.pkl'.format(save_name)), 'wb') as f: 38 | pickle.dump(log, f) 39 | if def_noise is not None: 40 | def_noise = (def_noise * 255).round() 41 | assert (def_noise.max()<=127 and def_noise.min()>=-128) 42 | def_noise = def_noise.astype(np.int8) 43 | with open(os.path.join(save_dir, '{}-def-noise.pkl'.format(save_name)), 'wb') as f: 44 | pickle.dump(def_noise, f) 45 | 46 | 47 | def main(args, logger): 48 | ''' init model / optim / dataloader / loss func ''' 49 | model = utils.get_arch(args.arch, args.dataset) 50 | optim = utils.get_optim( 51 | args.optim, model.parameters(), 52 | lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) 53 | train_loader = utils.get_indexed_loader( 54 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=True) 55 | test_loader = utils.get_indexed_loader( 56 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=False) 57 | criterion = torch.nn.CrossEntropyLoss() 58 | 59 | defender = attacks.PGDAttacker( 60 | radius = args.pgd_radius, 61 | steps = args.pgd_steps, 62 | step_size = args.pgd_step_size, 63 | random_start = args.pgd_random_start, 64 | norm_type = args.pgd_norm_type, 65 | ascending = False, 66 | ) 67 | 68 | if not args.cpu: 69 | model.cuda() 70 | criterion = criterion.cuda() 71 | 72 | if args.parallel: 73 | model = torch.nn.DataParallel(model) 74 | 75 | log = dict() 76 | 77 | ''' initialize the defensive noise (for unlearnable examples) ''' 78 | data_nums = len( train_loader.loader.dataset ) 79 | if args.dataset == 'cifar10' or args.dataset == 'cifar100': 80 | def_noise = np.zeros([data_nums, 3, 32, 32], dtype=np.float16) 81 | elif args.dataset == 'tiny-imagenet': 82 | def_noise = np.zeros([data_nums, 3, 64, 64], dtype=np.float16) 83 | elif args.dataset == 'imagenet-mini': 84 | def_noise = np.zeros([data_nums, 3, 224, 224], dtype=np.float16) 85 | else: 86 | raise NotImplementedError 87 | 88 | for step in range(args.train_steps): 89 | lr = args.lr * (args.lr_decay_rate ** (step // args.lr_decay_freq)) 90 | for group in optim.param_groups: 91 | group['lr'] = lr 92 | 93 | x, y, ii = next(train_loader) 94 | if not args.cpu: 95 | x, y = x.cuda(), y.cuda() 96 | 97 | if (step+1) % args.perturb_freq == 0: 98 | def_x = defender.perturb(model, criterion, x, y) 99 | def_noise[ii] = (def_x - x).cpu().numpy() 100 | 101 | if args.cpu: 102 | def_x = x + torch.tensor(def_noise[ii]) 103 | else: 104 | def_x = x + torch.tensor(def_noise[ii]).cuda() 105 | def_x.clamp_(-0.5, 0.5) 106 | 107 | model.train() 108 | _y = model(def_x) 109 | def_acc = (_y.argmax(dim=1) == y).sum().item() / len(x) 110 | def_loss = criterion(_y, y) 111 | optim.zero_grad() 112 | def_loss.backward() 113 | optim.step() 114 | 115 | utils.add_log(log, 'def_acc', def_acc) 116 | utils.add_log(log, 'def_loss', def_loss.item()) 117 | 118 | if (step+1) % args.save_freq == 0: 119 | save_checkpoint( 120 | args.save_dir, '{}-ckpt-{}'.format(args.save_name, step+1), 121 | model, optim, log, def_noise) 122 | 123 | if (step+1) % args.report_freq == 0: 124 | test_acc, test_loss = utils.evaluate(model, criterion, test_loader, args.cpu) 125 | utils.add_log(log, 'test_acc', test_acc) 126 | utils.add_log(log, 'test_loss', test_loss) 127 | 128 | logger.info('step [{}/{}]:'.format(step+1, args.train_steps)) 129 | logger.info('def_acc {:.2%} \t def_loss {:.3e}' 130 | .format( def_acc, def_loss.item() )) 131 | logger.info('test_acc {:.2%} \t test_loss {:.3e}' 132 | .format( test_acc, test_loss )) 133 | logger.info('') 134 | 135 | regenerate_def_noise( 136 | def_noise, model, criterion, train_loader, defender, args.cpu) 137 | 138 | save_checkpoint(args.save_dir, '{}-fin'.format(args.save_name), model, optim, log, def_noise) 139 | 140 | return 141 | 142 | 143 | if __name__ == '__main__': 144 | args = get_args() 145 | logger = utils.generic_init(args) 146 | 147 | try: 148 | main(args, logger) 149 | except Exception as e: 150 | logger.exception('Unexpected exception! %s', e) 151 | -------------------------------------------------------------------------------- /models/resnet.py: -------------------------------------------------------------------------------- 1 | ''' ref: 2 | https://github.com/kuangliu/pytorch-cifar/blob/master/models/resnet.py 3 | https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py 4 | ''' 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | 9 | 10 | class BasicBlock(nn.Module): 11 | expansion = 1 12 | 13 | def __init__(self, in_planes, planes, stride=1, wide=1): 14 | super(BasicBlock, self).__init__() 15 | planes = planes * wide 16 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 17 | self.bn1 = nn.BatchNorm2d(planes) 18 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) 19 | self.bn2 = nn.BatchNorm2d(planes) 20 | 21 | self.shortcut = nn.Sequential() 22 | if stride != 1 or in_planes != self.expansion*planes: 23 | self.shortcut = nn.Sequential( 24 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False), 25 | nn.BatchNorm2d(self.expansion*planes) 26 | ) 27 | 28 | def forward(self, x): 29 | out = F.relu(self.bn1(self.conv1(x))) 30 | out = self.bn2(self.conv2(out)) 31 | out += self.shortcut(x) 32 | out = F.relu(out) 33 | return out 34 | 35 | 36 | class Bottleneck(nn.Module): 37 | expansion = 4 38 | 39 | def __init__(self, in_planes, planes, stride=1, wide=1): 40 | super(Bottleneck, self).__init__() 41 | mid_planes = planes * wide 42 | self.conv1 = nn.Conv2d(in_planes, mid_planes, kernel_size=1, bias=False) 43 | self.bn1 = nn.BatchNorm2d(mid_planes) 44 | self.conv2 = nn.Conv2d(mid_planes, mid_planes, kernel_size=3, stride=stride, padding=1, bias=False) 45 | self.bn2 = nn.BatchNorm2d(mid_planes) 46 | self.conv3 = nn.Conv2d(mid_planes, self.expansion*planes, kernel_size=1, bias=False) 47 | self.bn3 = nn.BatchNorm2d(self.expansion*planes) 48 | 49 | self.shortcut = nn.Sequential() 50 | if stride != 1 or in_planes != self.expansion*planes: 51 | self.shortcut = nn.Sequential( 52 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False), 53 | nn.BatchNorm2d(self.expansion*planes) 54 | ) 55 | 56 | def forward(self, x): 57 | out = F.relu(self.bn1(self.conv1(x))) 58 | out = F.relu(self.bn2(self.conv2(out))) 59 | out = self.bn3(self.conv3(out)) 60 | out += self.shortcut(x) 61 | out = F.relu(out) 62 | return out 63 | 64 | 65 | class ResNet(nn.Module): 66 | def __init__(self, block, num_blocks, in_dims, out_dims, wide=1): 67 | super(ResNet, self).__init__() 68 | self.wide = wide 69 | self.in_planes = 64 70 | self.conv1 = nn.Conv2d(in_dims, 64, kernel_size=3, stride=1, padding=1, bias=False) 71 | self.bn1 = nn.BatchNorm2d(64) 72 | self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1) 73 | self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2) 74 | self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2) 75 | self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) 76 | self.avgpool = nn.AdaptiveAvgPool2d((1,1)) 77 | self.linear = nn.Linear(512*block.expansion, out_dims) 78 | 79 | def _make_layer(self, block, planes, num_blocks, stride): 80 | strides = [stride] + [1]*(num_blocks-1) 81 | layers = [] 82 | for stride in strides: 83 | layers.append(block(self.in_planes, planes, stride, self.wide)) 84 | self.in_planes = planes * block.expansion 85 | return nn.Sequential(*layers) 86 | 87 | def forward(self, x): 88 | out = F.relu(self.bn1(self.conv1(x))) 89 | out = self.layer1(out) 90 | out = self.layer2(out) 91 | out = self.layer3(out) 92 | out = self.layer4(out) 93 | out = self.avgpool(out) 94 | out = torch.flatten(out, 1) 95 | out = self.linear(out) 96 | return out 97 | 98 | def feature_extract(self, x): 99 | out = F.relu(self.bn1(self.conv1(x))) 100 | out = self.layer1(out) 101 | out = self.layer2(out) 102 | out = self.layer3(out) 103 | out = self.layer4(out) 104 | out = self.avgpool(out) 105 | out = torch.flatten(out, 1) 106 | return out 107 | 108 | 109 | class WRN(nn.Module): 110 | def __init__(self, num_blocks, in_dims, out_dims, wide=10): 111 | super(WRN, self).__init__() 112 | self.in_planes = 16 113 | self.wide = wide 114 | 115 | block = BasicBlock 116 | 117 | self.conv1 = nn.Conv2d(in_dims, self.in_planes, kernel_size=3, stride=1, padding=1, bias=False) 118 | self.bn1 = nn.BatchNorm2d(16) 119 | self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1) 120 | self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2) 121 | self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2) 122 | self.avgpool = nn.AdaptiveAvgPool2d((1,1)) 123 | self.linear = nn.Linear(64*wide, out_dims) 124 | 125 | def _make_layer(self, block, planes, num_blocks, stride): 126 | strides = [stride] + [1]*(num_blocks-1) 127 | layers = [] 128 | for stride in strides: 129 | layers.append(block(self.in_planes, planes, stride, self.wide)) 130 | self.in_planes = planes * self.wide * block.expansion 131 | 132 | return nn.Sequential(*layers) 133 | 134 | def forward(self, x): 135 | out = F.relu(self.bn1(self.conv1(x))) 136 | out = self.layer1(out) 137 | out = self.layer2(out) 138 | out = self.layer3(out) 139 | # out = F.avg_pool2d(out, 8) 140 | # out = out.view(out.shape[0], -1) 141 | out = self.avgpool(out) 142 | out = torch.flatten(out, 1) 143 | out = self.linear(out) 144 | return out 145 | 146 | 147 | def resnet18(in_dims, out_dims): 148 | return ResNet(BasicBlock, [2,2,2,2], in_dims, out_dims, 1) 149 | 150 | def wrn34_10(in_dims, out_dims): 151 | return WRN([5,5,5], in_dims, out_dims, wide=10) 152 | 153 | def resnet50(in_dims, out_dims): 154 | return ResNet(Bottleneck, [3,4,6,3], in_dims, out_dims, 1) 155 | -------------------------------------------------------------------------------- /generate_tap.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import argparse 4 | import numpy as np 5 | import torch 6 | 7 | import utils 8 | import attacks 9 | 10 | 11 | class TargetedIndexedDataset(): 12 | def __init__(self, dataset, classes): 13 | self.dataset = dataset 14 | self.classes = classes 15 | 16 | def __getitem__(self, idx): 17 | x, y, ii = self.dataset[idx] 18 | y += 1 19 | if y >= self.classes: y -= self.classes 20 | 21 | return x, y, ii 22 | 23 | def __len__(self): 24 | return len(self.dataset) 25 | 26 | 27 | def get_args(): 28 | parser = argparse.ArgumentParser() 29 | utils.add_shared_args(parser) 30 | 31 | parser.add_argument('--adv-type', type=str, default='robust-pgd', 32 | choices=['robust-pgd', 'diff-aug-pgd']) 33 | parser.add_argument('--targeted', action='store_true') 34 | 35 | parser.add_argument('--samp-num', type=int, default=1, 36 | help='set the number of samples for calculating expectations') 37 | 38 | parser.add_argument('--resume', action='store_true', 39 | help='set resume') 40 | parser.add_argument('--resume-path', type=str, default=None, 41 | help='set where to resume the model') 42 | 43 | return parser.parse_args() 44 | 45 | 46 | def regenerate_def_noise(def_noise, model, criterion, loader, defender, cpu, logger): 47 | cnt = 0 48 | for x, y, ii in loader: 49 | cnt += 1 50 | logger.info('progress [{}/{}]'.format(cnt, len(loader)) ) 51 | 52 | if not cpu: x, y = x.cuda(), y.cuda() 53 | delta = defender.perturb(model, criterion, x, y) 54 | def_noise[ii] = delta.cpu().numpy() 55 | 56 | 57 | def save_checkpoint(save_dir, save_name, model, optim, log, def_noise=None): 58 | torch.save({ 59 | 'model_state_dict': utils.get_model_state(model), 60 | 'optim_state_dict': optim.state_dict(), 61 | }, os.path.join(save_dir, '{}-model.pkl'.format(save_name))) 62 | with open(os.path.join(save_dir, '{}-log.pkl'.format(save_name)), 'wb') as f: 63 | pickle.dump(log, f) 64 | if def_noise is not None: 65 | def_noise = (def_noise * 255).round() 66 | assert (def_noise.max()<=127 and def_noise.min()>=-128) 67 | def_noise = def_noise.astype(np.int8) 68 | with open(os.path.join(save_dir, '{}-def-noise.pkl'.format(save_name)), 'wb') as f: 69 | pickle.dump(def_noise, f) 70 | 71 | 72 | def main(args, logger): 73 | ''' init model / optim / loss func ''' 74 | model = utils.get_arch(args.arch, args.dataset) 75 | optim = utils.get_optim( 76 | args.optim, model.parameters(), 77 | lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) 78 | criterion = torch.nn.CrossEntropyLoss() 79 | 80 | ''' get Tensor train loader ''' 81 | train_loader = utils.get_indexed_tensor_loader( 82 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=True) 83 | 84 | dataset = train_loader.loader.dataset 85 | ascending = True 86 | if args.targeted: 87 | if args.dataset == 'cifar10': classes = 10 88 | elif args.dataset == 'cifar100': classes = 100 89 | elif args.dataset == 'imagenet-mini': classes = 100 90 | else: raise ValueError 91 | dataset = TargetedIndexedDataset(dataset, classes) 92 | ascending = False 93 | 94 | train_loader = utils.Loader(dataset, batch_size=args.batch_size, shuffle=False, drop_last=False) 95 | 96 | ''' get train transforms ''' 97 | train_trans = utils.get_transforms( 98 | args.dataset, train=True, is_tensor=True) 99 | 100 | ''' get (original) test loader ''' 101 | test_loader = utils.get_indexed_loader( 102 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=False) 103 | 104 | if args.adv_type == 'robust-pgd': 105 | defender = attacks.RobustPGDAttacker( 106 | samp_num = args.samp_num, 107 | trans = train_trans, 108 | radius = args.pgd_radius, 109 | steps = args.pgd_steps, 110 | step_size = args.pgd_step_size, 111 | random_start = args.pgd_random_start, 112 | ascending = ascending, 113 | ) 114 | elif args.adv_type == 'diff-aug-pgd': 115 | defender = attacks.DiffAugPGDAttacker( 116 | samp_num = args.samp_num, 117 | trans = train_trans, 118 | radius = args.pgd_radius, 119 | steps = args.pgd_steps, 120 | step_size = args.pgd_step_size, 121 | random_start = args.pgd_random_start, 122 | ascending = ascending, 123 | ) 124 | else: raise ValueError 125 | 126 | ''' initialize the defensive noise (for unlearnable examples) ''' 127 | data_nums = len( train_loader.loader.dataset ) 128 | if args.dataset == 'cifar10' or args.dataset == 'cifar100': 129 | def_noise = np.zeros([data_nums, 3, 32, 32], dtype=np.float16) 130 | elif args.dataset == 'tiny-imagenet': 131 | def_noise = np.zeros([data_nums, 3, 64, 64], dtype=np.float16) 132 | elif args.dataset == 'imagenet-mini': 133 | def_noise = np.zeros([data_nums, 3, 256, 256], dtype=np.float16) 134 | else: 135 | raise NotImplementedError 136 | 137 | start_step = 0 138 | log = dict() 139 | 140 | if not args.cpu: 141 | model.cuda() 142 | criterion = criterion.cuda() 143 | 144 | if args.resume: 145 | state_dict = torch.load( os.path.join(args.resume_path) ) 146 | model.load_state_dict( state_dict['model_state_dict'] ) 147 | optim.load_state_dict( state_dict['optim_state_dict'] ) 148 | del state_dict 149 | 150 | if args.parallel: 151 | model = torch.nn.DataParallel(model) 152 | 153 | logger.info('Noise generation started') 154 | 155 | regenerate_def_noise( 156 | def_noise, model, criterion, train_loader, defender, args.cpu, logger) 157 | 158 | logger.info('Noise generation finished') 159 | 160 | save_checkpoint(args.save_dir, '{}-fin'.format(args.save_name), model, optim, log, def_noise) 161 | 162 | return 163 | 164 | 165 | if __name__ == '__main__': 166 | args = get_args() 167 | logger = utils.generic_init(args) 168 | 169 | logger.info('EXP: robust minimax pgd perturbation') 170 | try: 171 | main(args, logger) 172 | except Exception as e: 173 | logger.exception('Unexpected exception! %s', e) 174 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import argparse 4 | import numpy as np 5 | import torch 6 | 7 | import utils 8 | import attacks 9 | 10 | 11 | def get_args(): 12 | parser = argparse.ArgumentParser() 13 | utils.add_shared_args(parser) 14 | 15 | parser.add_argument('--data-mode', type=str, default='mix', 16 | choices=['mix', 'clear', 'manual'], 17 | help='mix = clear + unlearnable data, clear = clear data only') 18 | 19 | parser.add_argument('--filter', type=str, default=None, 20 | choices=['averaging', 'gaussian', 'median', 'bilateral'], 21 | help='select the low pass filter; only works in [mix] mode') 22 | 23 | parser.add_argument('--man-data-path', type=str, default=None, 24 | help='set the path to the manual dataset') 25 | parser.add_argument('--noise-path', type=str, default=None, 26 | help='set the path to the train images noises') 27 | parser.add_argument('--poi-idx-path', type=str, default=None, 28 | help='set the path to the poisoned indices') 29 | parser.add_argument('--resume-path', type=str, default=None, 30 | help='set where to resume the model') 31 | 32 | parser.add_argument('--perturb-freq', type=int, default=1, 33 | help='set the perturbation frequency') 34 | parser.add_argument('--report-freq', type=int, default=500, 35 | help='set the report frequency') 36 | parser.add_argument('--save-freq', type=int, default=5000, 37 | help='set the checkpoint saving frequency') 38 | 39 | return parser.parse_args() 40 | 41 | 42 | def get_manual_loader(dataset, data_path, batch_size): 43 | with open(data_path, 'rb') as f: 44 | man_dataset = pickle.load(f) 45 | trans = utils.get_transforms(dataset, train=True, is_tensor=False) 46 | 47 | man_dataset = utils.Dataset( man_dataset['x'], man_dataset['y'].astype(np.int64), trans ) 48 | loader = utils.Loader(man_dataset, batch_size=batch_size, shuffle=True, drop_last=True) 49 | 50 | return loader 51 | 52 | 53 | def save_checkpoint(save_dir, save_name, model, optim, log): 54 | torch.save({ 55 | 'model_state_dict': utils.get_model_state(model), 56 | 'optim_state_dict': optim.state_dict(), 57 | }, os.path.join(save_dir, '{}-model.pkl'.format(save_name))) 58 | with open(os.path.join(save_dir, '{}-log.pkl'.format(save_name)), 'wb') as f: 59 | pickle.dump(log, f) 60 | 61 | 62 | def main(args, logger): 63 | ''' init model / optim / dataloader / loss func ''' 64 | model = utils.get_arch(args.arch, args.dataset) 65 | 66 | if args.resume_path is not None: 67 | state_dict = torch.load(args.resume_path, map_location=torch.device('cpu')) 68 | model.load_state_dict( state_dict['model_state_dict'] ) 69 | del state_dict 70 | 71 | criterion = torch.nn.CrossEntropyLoss() 72 | 73 | if not args.cpu: 74 | model.cuda() 75 | criterion = criterion.cuda() 76 | 77 | if args.parallel: 78 | model = torch.nn.DataParallel(model) 79 | 80 | optim = utils.get_optim( 81 | args.optim, model.parameters(), 82 | lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) 83 | 84 | if args.data_mode == 'mix': 85 | train_loader = utils.get_poisoned_loader( 86 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=True, 87 | noise_path=args.noise_path, poisoned_indices_path=args.poi_idx_path, fitr=args.filter) 88 | elif args.data_mode == 'clear': 89 | train_loader = utils.get_clear_loader( 90 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=True, 91 | poisoned_indices_path=args.poi_idx_path, fitr=args.filter) 92 | elif args.data_mode == 'manual': 93 | train_loader = get_manual_loader(args.dataset, args.man_data_path, args.batch_size) 94 | else: 95 | raise NotImplementedError 96 | 97 | test_loader = utils.get_poisoned_loader( 98 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=False) 99 | 100 | attacker = attacks.PGDAttacker( 101 | radius = args.pgd_radius, 102 | steps = args.pgd_steps, 103 | step_size = args.pgd_step_size, 104 | random_start = args.pgd_random_start, 105 | norm_type = args.pgd_norm_type, 106 | ascending = True, 107 | ) 108 | 109 | log = dict() 110 | 111 | for step in range(args.train_steps): 112 | lr = args.lr * (args.lr_decay_rate ** (step // args.lr_decay_freq)) 113 | for group in optim.param_groups: 114 | group['lr'] = lr 115 | 116 | x, y = next(train_loader) 117 | if not args.cpu: 118 | x, y = x.cuda(), y.cuda() 119 | 120 | if (step+1) % args.perturb_freq == 0: 121 | adv_x = attacker.perturb(model, criterion, x, y) 122 | else: 123 | adv_x = x 124 | 125 | model.train() 126 | _y = model(adv_x) 127 | adv_acc = (_y.argmax(dim=1) == y).sum().item() / len(x) 128 | adv_loss = criterion(_y, y) 129 | optim.zero_grad() 130 | adv_loss.backward() 131 | optim.step() 132 | 133 | utils.add_log(log, 'adv_acc', adv_acc) 134 | utils.add_log(log, 'adv_loss', adv_loss.item()) 135 | 136 | if (step+1) % args.save_freq == 0: 137 | save_checkpoint( 138 | args.save_dir, '{}-ckpt-{}'.format(args.save_name, step+1), 139 | model, optim, log) 140 | 141 | if (step+1) % args.report_freq == 0: 142 | test_acc, test_loss = utils.evaluate(model, criterion, test_loader, args.cpu) 143 | utils.add_log(log, 'test_acc', test_acc) 144 | utils.add_log(log, 'test_loss', test_loss) 145 | 146 | logger.info('step [{}/{}]:'.format(step+1, args.train_steps)) 147 | logger.info('adv_acc {:.2%} \t adv_loss {:.3e}' 148 | .format( adv_acc, adv_loss.item() )) 149 | logger.info('test_acc {:.2%} \t test_loss {:.3e}' 150 | .format( test_acc, test_loss )) 151 | logger.info('') 152 | 153 | save_checkpoint(args.save_dir, '{}-fin'.format(args.save_name), model, optim, log) 154 | 155 | return 156 | 157 | 158 | if __name__ == '__main__': 159 | args = get_args() 160 | logger = utils.generic_init(args) 161 | 162 | try: 163 | main(args, logger) 164 | except Exception as e: 165 | logger.exception('Unexpected exception! %s', e) 166 | -------------------------------------------------------------------------------- /utils/imagenet_utils.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import numpy as np 3 | import torch 4 | import cv2 5 | import torchvision 6 | import os 7 | import pickle 8 | import torchvision.transforms as transforms 9 | 10 | 11 | class ElementWiseTransform(): 12 | def __init__(self, trans=None): 13 | self.trans = trans 14 | 15 | def __call__(self, x): 16 | if self.trans is None: return x 17 | return torch.cat( [self.trans( xx.view(1, *xx.shape) ) for xx in x] ) 18 | 19 | 20 | class Dataset(torchvision.datasets.DatasetFolder): 21 | def __init__(self, dataset): 22 | assert isinstance(dataset, torchvision.datasets.DatasetFolder) 23 | self.loader = dataset.loader 24 | self.classes = dataset.classes 25 | self.samples = dataset.samples 26 | self.y = dataset.targets 27 | self.transform = dataset.transform 28 | self.target_transform = None 29 | 30 | def __getitem__(self, idx): 31 | return super().__getitem__(idx) 32 | 33 | def __len__(self): 34 | return len(self.y) 35 | 36 | 37 | class IndexedTensorDataset(Dataset): 38 | def __init__(self, dataset): 39 | assert isinstance(dataset, Dataset) 40 | self.loader = dataset.loader 41 | self.classes = dataset.classes 42 | self.samples = dataset.samples 43 | self.y = dataset.y 44 | self.transform = transforms.Compose([ transforms.Resize( [256, 256] ) ]) 45 | self.target_transform = None 46 | 47 | def __getitem__(self, idx): 48 | x, y = super().__getitem__(idx) 49 | ''' transform HWC pic to CWH pic ''' 50 | x = np.asarray(x, dtype=np.uint8) 51 | x = torch.tensor(x, dtype=torch.float32).permute(2,0,1) 52 | return x, y, idx 53 | 54 | def __len__(self): 55 | return len(self.y) 56 | 57 | 58 | class PoisonedDataset(Dataset): 59 | def __init__(self, dataset, noise, fitr=None): 60 | assert isinstance(dataset, Dataset) 61 | self.loader = dataset.loader 62 | self.classes = dataset.classes 63 | self.samples = dataset.samples 64 | self.y = dataset.y 65 | self.transform = transforms.Compose([ transforms.Resize([ noise.shape[1], noise.shape[2] ]) ]) 66 | self.target_transform = None 67 | self.data_transform = dataset.transform 68 | self.data_fitr = fitr 69 | ''' the shape of the noise should be (NHWC) ''' 70 | self.noise = noise 71 | 72 | def __getitem__(self, idx): 73 | x, y = super().__getitem__(idx) 74 | x = (np.asarray(x, dtype=np.int16) + self.noise[idx].astype(np.int16)).clip(0, 255).astype(np.uint8) 75 | 76 | ''' low pass filtering ''' 77 | if self.data_fitr is not None: 78 | x = self.data_fitr(x) 79 | 80 | x = self.data_transform( Image.fromarray(x) ) 81 | # x = self.data_transform( Image.fromarray(x.astype(np.uint8)) ) 82 | return x, y 83 | 84 | def __len__(self): 85 | return len(self.y) 86 | 87 | 88 | class IndexedDataset(Dataset): 89 | def __init__(self, dataset): 90 | assert isinstance(dataset, Dataset) 91 | self.loader = dataset.loader 92 | self.classes = dataset.classes 93 | self.samples = dataset.samples 94 | self.y = dataset.y 95 | self.transform = dataset.transform 96 | self.target_transform = None 97 | 98 | def __getitem__(self, idx): 99 | x, y = super().__getitem__(idx) 100 | return x, y, idx 101 | 102 | def __len__(self): 103 | return len(self.y) 104 | 105 | 106 | class Loader(): 107 | def __init__(self, dataset, batch_size, shuffle=False, drop_last=False, num_workers=4): 108 | self.loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers) 109 | self.iterator = None 110 | 111 | def __iter__(self): 112 | return iter(self.loader) 113 | 114 | def __len__(self): 115 | return len(self.loader) 116 | 117 | def __next__(self): 118 | if self.iterator is None: 119 | self.iterator = iter(self.loader) 120 | 121 | try: 122 | samples = next(self.iterator) 123 | except StopIteration: 124 | self.iterator = iter(self.loader) 125 | samples = next(self.iterator) 126 | 127 | return samples 128 | 129 | 130 | def datasetImageNet(root='./data', train=True, transform=None): 131 | if train: root = os.path.join(root, 'ILSVRC2012_img_train') 132 | else: root = os.path.join(root, 'ILSVRC2012_img_val') 133 | return torchvision.datasets.ImageFolder(root=root, transform=transform) 134 | 135 | 136 | def datasetImageNetMini(root='./data', train=True, transform=None): 137 | dataset = datasetImageNet(root=root, train=train, transform=transform) 138 | ''' imagenet-mini is a subset of the first 100 classes of ImageNet ''' 139 | idx = np.where( np.array(dataset.targets) < 100 )[0] 140 | dataset.samples = [ dataset.samples[ii] for ii in idx ] 141 | dataset.targets = [ dataset.targets[ii] for ii in idx ] 142 | return dataset 143 | 144 | 145 | def get_transforms(dataset, train=True, is_tensor=True): 146 | assert (dataset == 'imagenet' or dataset == 'imagenet-mini') 147 | if train: 148 | comp1 = [ 149 | transforms.RandomResizedCrop(224), 150 | transforms.RandomHorizontalFlip(), ] 151 | else: 152 | comp1 = [ 153 | transforms.Resize( [256, 256] ), 154 | transforms.CenterCrop(224), ] 155 | 156 | if is_tensor: 157 | comp2 = [ 158 | torchvision.transforms.Normalize((255*0.5, 255*0.5, 255*0.5), (255., 255., 255.))] 159 | else: 160 | comp2 = [ 161 | transforms.ToTensor(), 162 | transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.))] 163 | 164 | trans = transforms.Compose( [*comp1, *comp2] ) 165 | 166 | if is_tensor: trans = ElementWiseTransform(trans) 167 | 168 | return trans 169 | 170 | 171 | def get_filter(fitr): 172 | if fitr == 'averaging': 173 | return lambda x: cv2.blur(x, (3,3)) 174 | elif fitr == 'gaussian': 175 | return lambda x: cv2.GaussianBlur(x, (3,3), 0) 176 | elif fitr == 'median': 177 | return lambda x: cv2.medianBlur(x, 3) 178 | elif fitr == 'bilateral': 179 | return lambda x: cv2.bilateralFilter(x, 9, 75, 75) 180 | 181 | raise ValueError 182 | 183 | 184 | def get_dataset(dataset, root='./data', train=True): 185 | assert (dataset == 'imagenet' or dataset == 'imagenet-mini') 186 | 187 | transform = get_transforms(dataset, train=train, is_tensor=False) 188 | 189 | if dataset == 'imagenet': 190 | target_set = datasetImageNet(root=root, train=train, transform=transform) 191 | if dataset == 'imagenet-mini': 192 | target_set = datasetImageNetMini(root=root, train=train, transform=transform) 193 | 194 | target_set = Dataset(target_set) 195 | 196 | return target_set 197 | 198 | 199 | def get_indexed_loader(dataset, batch_size, root='./data', train=True): 200 | target_set = get_dataset(dataset, root=root, train=train) 201 | 202 | if train: 203 | target_set = IndexedDataset(target_set) 204 | else: 205 | pass 206 | 207 | if train: 208 | loader = Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 209 | else: 210 | loader = Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 211 | 212 | return loader 213 | 214 | 215 | def get_indexed_tensor_loader(dataset, batch_size, root='./data', train=True): 216 | target_set = get_dataset(dataset, root=root, train=train) 217 | target_set = IndexedTensorDataset(target_set) 218 | 219 | if train: 220 | loader = Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 221 | else: 222 | loader = Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 223 | 224 | return loader 225 | 226 | 227 | def get_poisoned_loader( 228 | dataset, batch_size, root='./data', train=True, 229 | noise_path=None, noise_rate=1.0, poisoned_indices_path=None, fitr=None): 230 | 231 | target_set = get_dataset(dataset, root=root, train=train) 232 | 233 | if noise_path is not None: 234 | with open(noise_path, 'rb') as f: 235 | raw_noise = pickle.load(f) 236 | 237 | assert isinstance(raw_noise, np.ndarray) 238 | assert raw_noise.dtype == np.int8 239 | 240 | # raw_noise = raw_noise.astype(np.int32) 241 | 242 | if noise_rate < 1.0: 243 | raise NotImplementedError 244 | 245 | noise = raw_noise 246 | 247 | ''' restore noise (NCWH) for raw images (NWHC) ''' 248 | noise = np.transpose(noise, [0,2,3,1]) 249 | 250 | lp_fitr = None if fitr is None else get_filter(fitr) 251 | 252 | target_set = PoisonedDataset(target_set, noise, fitr=lp_fitr) 253 | 254 | else: 255 | pass 256 | 257 | if train: 258 | loader = Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 259 | else: 260 | loader = Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 261 | 262 | return loader 263 | 264 | 265 | def get_clear_loader( 266 | dataset, batch_size, root='./data', train=True, 267 | noise_rate=1.0, poisoned_indices_path=None): 268 | 269 | raise NotImplementedError 270 | -------------------------------------------------------------------------------- /generate_robust_em.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import argparse 4 | import numpy as np 5 | import torch 6 | 7 | import utils 8 | import attacks 9 | 10 | 11 | def get_args(): 12 | parser = argparse.ArgumentParser() 13 | utils.add_shared_args(parser) 14 | 15 | parser.add_argument('--perturb-freq', type=int, default=1, 16 | help='set the perturbation frequency') 17 | parser.add_argument('--report-freq', type=int, default=500, 18 | help='set the report frequency') 19 | parser.add_argument('--save-freq', type=int, default=5000, 20 | help='set the checkpoint saving frequency') 21 | 22 | parser.add_argument('--samp-num', type=int, default=1, 23 | help='set the number of samples for calculating expectations') 24 | 25 | parser.add_argument('--atk-pgd-radius', type=float, default=0, 26 | help='set the adv perturbation radius in minimax-pgd') 27 | parser.add_argument('--atk-pgd-steps', type=int, default=0, 28 | help='set the number of adv iteration steps in minimax-pgd') 29 | parser.add_argument('--atk-pgd-step-size', type=float, default=0, 30 | help='set the adv step size in minimax-pgd') 31 | parser.add_argument('--atk-pgd-random-start', action='store_true', 32 | help='if select, randomly choose starting points each time performing adv pgd in minimax-pgd') 33 | 34 | parser.add_argument('--pretrain', action='store_true', 35 | help='if select, use pre-trained model') 36 | parser.add_argument('--pretrain-path', type=str, default=None, 37 | help='set the path to the pretrained model') 38 | 39 | parser.add_argument('--resume', action='store_true', 40 | help='set resume') 41 | parser.add_argument('--resume-step', type=int, default=None, 42 | help='set which step to resume the model') 43 | parser.add_argument('--resume-dir', type=str, default=None, 44 | help='set where to resume the model') 45 | parser.add_argument('--resume-name', type=str, default=None, 46 | help='set the resume name') 47 | 48 | return parser.parse_args() 49 | 50 | 51 | def load_pretrained_model(model, arch, pre_state_dict): 52 | target_state_dict = model.state_dict() 53 | 54 | for name, param in pre_state_dict.items(): 55 | if (arch=='resnet18') and ('linear' in name): continue 56 | target_state_dict[name].copy_(param) 57 | 58 | 59 | def regenerate_def_noise(def_noise, model, criterion, loader, defender, cpu): 60 | for x, y, ii in loader: 61 | if not cpu: x, y = x.cuda(), y.cuda() 62 | delta = defender.perturb(model, criterion, x, y) 63 | # def_noise[ii] = delta.cpu().numpy() 64 | def_noise[ii] = (delta.cpu().numpy() * 255).round().astype(np.int8) 65 | 66 | 67 | def save_checkpoint(save_dir, save_name, model, optim, log, def_noise=None): 68 | torch.save({ 69 | 'model_state_dict': utils.get_model_state(model), 70 | 'optim_state_dict': optim.state_dict(), 71 | }, os.path.join(save_dir, '{}-model.pkl'.format(save_name))) 72 | with open(os.path.join(save_dir, '{}-log.pkl'.format(save_name)), 'wb') as f: 73 | pickle.dump(log, f) 74 | if def_noise is not None: 75 | # def_noise = (def_noise * 255).round() 76 | # def_noise = def_noise * 255 77 | # def_noise = def_noise.round() 78 | # assert (def_noise.max()<=127 and def_noise.min()>=-128) 79 | # def_noise = def_noise.astype(np.int8) 80 | with open(os.path.join(save_dir, '{}-def-noise.pkl'.format(save_name)), 'wb') as f: 81 | pickle.dump(def_noise, f) 82 | 83 | 84 | def main(args, logger): 85 | ''' init model / optim / loss func ''' 86 | model = utils.get_arch(args.arch, args.dataset) 87 | optim = utils.get_optim( 88 | args.optim, model.parameters(), 89 | lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) 90 | criterion = torch.nn.CrossEntropyLoss() 91 | 92 | ''' get Tensor train loader ''' 93 | # trainset = utils.get_dataset(args.dataset, root=args.data_dir, train=True) 94 | # trainset = utils.IndexedTensorDataset(trainset.x, trainset.y) 95 | # train_loader = utils.Loader( 96 | # trainset, batch_size=args.batch_size, shuffle=True, drop_last=True) 97 | train_loader = utils.get_indexed_tensor_loader( 98 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=True) 99 | 100 | ''' get train transforms ''' 101 | train_trans = utils.get_transforms( 102 | args.dataset, train=True, is_tensor=True) 103 | 104 | ''' get (original) test loader ''' 105 | test_loader = utils.get_indexed_loader( 106 | args.dataset, batch_size=args.batch_size, root=args.data_dir, train=False) 107 | 108 | defender = attacks.RobustMinimaxPGDDefender( 109 | samp_num = args.samp_num, 110 | trans = train_trans, 111 | radius = args.pgd_radius, 112 | steps = args.pgd_steps, 113 | step_size = args.pgd_step_size, 114 | random_start = args.pgd_random_start, 115 | atk_radius = args.atk_pgd_radius, 116 | atk_steps = args.atk_pgd_steps, 117 | atk_step_size = args.atk_pgd_step_size, 118 | atk_random_start = args.atk_pgd_random_start, 119 | ) 120 | 121 | attacker = attacks.PGDAttacker( 122 | radius = args.atk_pgd_radius, 123 | steps = args.atk_pgd_steps, 124 | step_size = args.atk_pgd_step_size, 125 | random_start = args.atk_pgd_random_start, 126 | norm_type = 'l-infty', 127 | ascending = True, 128 | ) 129 | 130 | ''' initialize the defensive noise (for unlearnable examples) ''' 131 | data_nums = len( train_loader.loader.dataset ) 132 | if args.dataset == 'cifar10' or args.dataset == 'cifar100': 133 | # def_noise = np.zeros([data_nums, 3, 32, 32], dtype=np.float16) 134 | def_noise = np.zeros([data_nums, 3, 32, 32], dtype=np.int8) 135 | elif args.dataset == 'tiny-imagenet': 136 | # def_noise = np.zeros([data_nums, 3, 64, 64], dtype=np.float16) 137 | def_noise = np.zeros([data_nums, 3, 64, 64], dtype=np.int8) 138 | elif args.dataset == 'imagenet-mini': 139 | # def_noise = np.zeros([data_nums, 3, 256, 256], dtype=np.float16) 140 | def_noise = np.zeros([data_nums, 3, 256, 256], dtype=np.int8) 141 | else: 142 | raise NotImplementedError 143 | 144 | start_step = 0 145 | log = dict() 146 | 147 | if not args.cpu: 148 | model.cuda() 149 | criterion = criterion.cuda() 150 | 151 | if args.pretrain: 152 | state_dict = torch.load(args.pretrain_path) 153 | load_pretrained_model(model, args.arch, state_dict['model_state_dict']) 154 | del state_dict 155 | 156 | if args.resume: 157 | start_step = args.resume_step 158 | 159 | state_dict = torch.load( os.path.join(args.resume_dir, '{}-model.pkl'.format(args.resume_name)) ) 160 | model.load_state_dict( state_dict['model_state_dict'] ) 161 | optim.load_state_dict( state_dict['optim_state_dict'] ) 162 | del state_dict 163 | 164 | with open(os.path.join(args.resume_dir, '{}-log.pkl'.format(args.resume_name)), 'rb') as f: 165 | log = pickle.load(f) 166 | 167 | with open(os.path.join(args.resume_dir, '{}-def-noise.pkl'.format(args.resume_name)), 'rb') as f: 168 | # def_noise = pickle.load(f).astype(np.float16) / 255 169 | def_noise = pickle.load(f) 170 | 171 | if args.parallel: 172 | model = torch.nn.DataParallel(model) 173 | 174 | for step in range(start_step, args.train_steps): 175 | lr = args.lr * (args.lr_decay_rate ** (step // args.lr_decay_freq)) 176 | for group in optim.param_groups: 177 | group['lr'] = lr 178 | 179 | x, y, ii = next(train_loader) 180 | if not args.cpu: 181 | x, y = x.cuda(), y.cuda() 182 | 183 | if (step+1) % args.perturb_freq == 0: 184 | delta = defender.perturb(model, criterion, x, y) 185 | # def_noise[ii] = delta.cpu().numpy() 186 | def_noise[ii] = (delta.cpu().numpy() * 255).round().astype(np.int8) 187 | 188 | if args.cpu: 189 | # def_x = train_trans(x + torch.tensor(def_noise[ii])*255) 190 | def_x = train_trans(x + torch.tensor(def_noise[ii])) 191 | else: 192 | # def_x = train_trans(x + torch.tensor(def_noise[ii]).cuda()*255) 193 | def_x = train_trans(x + torch.tensor(def_noise[ii]).cuda()) 194 | def_x.clamp_(-0.5, 0.5) 195 | 196 | adv_x = attacker.perturb(model, criterion, def_x, y) 197 | 198 | model.train() 199 | _y = model(adv_x) 200 | def_acc = (_y.argmax(dim=1) == y).sum().item() / len(x) 201 | def_loss = criterion(_y, y) 202 | optim.zero_grad() 203 | def_loss.backward() 204 | optim.step() 205 | 206 | utils.add_log(log, 'def_acc', def_acc) 207 | utils.add_log(log, 'def_loss', def_loss.item()) 208 | 209 | if (step+1) % args.save_freq == 0: 210 | save_checkpoint( 211 | args.save_dir, '{}-ckpt-{}'.format(args.save_name, step+1), 212 | model, optim, log, def_noise) 213 | 214 | if (step+1) % args.report_freq == 0: 215 | test_acc, test_loss = utils.evaluate(model, criterion, test_loader, args.cpu) 216 | utils.add_log(log, 'test_acc', test_acc) 217 | utils.add_log(log, 'test_loss', test_loss) 218 | 219 | logger.info('step [{}/{}]:'.format(step+1, args.train_steps)) 220 | logger.info('def_acc {:.2%} \t def_loss {:.3e}' 221 | .format( def_acc, def_loss.item() )) 222 | logger.info('test_acc {:.2%} \t test_loss {:.3e}' 223 | .format( test_acc, test_loss )) 224 | logger.info('') 225 | 226 | logger.info('Noise generation started') 227 | 228 | regenerate_def_noise( 229 | def_noise, model, criterion, train_loader, defender, args.cpu) 230 | 231 | logger.info('Noise generation finished') 232 | 233 | save_checkpoint(args.save_dir, '{}-fin'.format(args.save_name), model, optim, log, def_noise) 234 | 235 | return 236 | 237 | 238 | if __name__ == '__main__': 239 | args = get_args() 240 | logger = utils.generic_init(args) 241 | 242 | logger.info('EXP: robust minimax pgd perturbation') 243 | try: 244 | main(args, logger) 245 | except Exception as e: 246 | logger.exception('Unexpected exception! %s', e) 247 | -------------------------------------------------------------------------------- /utils/generic.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import os 3 | import sys 4 | import logging 5 | import numpy as np 6 | import torch 7 | import cv2 8 | import torchvision 9 | import torchvision.transforms as transforms 10 | 11 | import models 12 | from . import data 13 | from . import imagenet_utils 14 | 15 | 16 | class AverageMeter(): 17 | def __init__(self): 18 | self.cnt = 0 19 | self.sum = 0 20 | self.mean = 0 21 | 22 | def update(self, val, cnt): 23 | self.cnt += cnt 24 | self.sum += val * cnt 25 | self.mean = self.sum / self.cnt 26 | 27 | def average(self): 28 | return self.mean 29 | 30 | def total(self): 31 | return self.sum 32 | 33 | 34 | def add_log(log, key, value): 35 | if key not in log.keys(): 36 | log[key] = [] 37 | log[key].append(value) 38 | 39 | 40 | def get_transforms(dataset, train=True, is_tensor=True): 41 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 42 | return imagenet_utils.get_transforms(dataset, train, is_tensor) 43 | 44 | if train: 45 | if dataset == 'cifar10' or dataset == 'cifar100': 46 | comp1 = [ 47 | transforms.RandomHorizontalFlip(), 48 | transforms.RandomCrop(32, 4), ] 49 | elif dataset == 'tiny-imagenet': 50 | comp1 = [ 51 | transforms.RandomHorizontalFlip(), 52 | transforms.RandomCrop(64, 8), ] 53 | else: 54 | raise NotImplementedError 55 | else: 56 | comp1 = [] 57 | 58 | if is_tensor: 59 | comp2 = [ 60 | torchvision.transforms.Normalize((255*0.5, 255*0.5, 255*0.5), (255., 255., 255.))] 61 | else: 62 | comp2 = [ 63 | transforms.ToTensor(), 64 | transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.))] 65 | 66 | trans = transforms.Compose( [*comp1, *comp2] ) 67 | 68 | if is_tensor: trans = data.ElementWiseTransform(trans) 69 | 70 | return trans 71 | 72 | 73 | def get_filter(fitr): 74 | if fitr == 'averaging': 75 | return lambda x: cv2.blur(x, (3,3)) 76 | elif fitr == 'gaussian': 77 | return lambda x: cv2.GaussianBlur(x, (3,3), 0) 78 | elif fitr == 'median': 79 | return lambda x: cv2.medianBlur(x, 3) 80 | elif fitr == 'bilateral': 81 | return lambda x: cv2.bilateralFilter(x, 9, 75, 75) 82 | 83 | raise ValueError 84 | 85 | 86 | def get_dataset(dataset, root='./data', train=True, fitr=None): 87 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 88 | return imagenet_utils.get_dataset(dataset, root, train) 89 | 90 | transform = get_transforms(dataset, train=train, is_tensor=False) 91 | lp_fitr = None if fitr is None else get_filter(fitr) 92 | 93 | if dataset == 'cifar10': 94 | target_set = data.datasetCIFAR10(root=root, train=train, transform=transform) 95 | x, y = target_set.data, target_set.targets 96 | elif dataset == 'cifar100': 97 | target_set = data.datasetCIFAR100(root=root, train=train, transform=transform) 98 | x, y = target_set.data, target_set.targets 99 | elif dataset == 'tiny-imagenet': 100 | target_set = data.datasetTinyImageNet(root=root, train=train, transform=transform) 101 | x, y = target_set.x, target_set.y 102 | else: 103 | raise NotImplementedError('dataset {} is not supported'.format(dataset)) 104 | 105 | return data.Dataset(x, y, transform, lp_fitr) 106 | 107 | 108 | def get_indexed_loader(dataset, batch_size, root='./data', train=True): 109 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 110 | return imagenet_utils.get_indexed_loader(dataset, batch_size, root, train) 111 | 112 | target_set = get_dataset(dataset, root=root, train=train) 113 | 114 | if train: 115 | target_set = data.IndexedDataset(x=target_set.x, y=target_set.y, transform=target_set.transform) 116 | else: 117 | target_set = data.Dataset(x=target_set.x, y=target_set.y, transform=target_set.transform) 118 | 119 | if train: 120 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 121 | else: 122 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 123 | 124 | return loader 125 | 126 | 127 | def get_indexed_tensor_loader(dataset, batch_size, root='./data', train=True): 128 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 129 | return imagenet_utils.get_indexed_tensor_loader(dataset, batch_size, root, train) 130 | 131 | target_set = get_dataset(dataset, root=root, train=train) 132 | target_set = data.IndexedTensorDataset(x=target_set.x, y=target_set.y) 133 | 134 | if train: 135 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 136 | else: 137 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 138 | 139 | return loader 140 | 141 | 142 | def get_poisoned_loader( 143 | dataset, batch_size, root='./data', train=True, 144 | noise_path=None, noise_rate=1.0, poisoned_indices_path=None, fitr=None): 145 | 146 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 147 | return imagenet_utils.get_poisoned_loader( 148 | dataset, batch_size, root, train, noise_path, noise_rate, poisoned_indices_path, fitr) 149 | 150 | target_set = get_dataset(dataset, root=root, train=train, fitr=fitr) 151 | 152 | if noise_path is not None: 153 | with open(noise_path, 'rb') as f: 154 | raw_noise = pickle.load(f) 155 | 156 | assert isinstance(raw_noise, np.ndarray) 157 | assert raw_noise.dtype == np.int8 158 | 159 | raw_noise = raw_noise.astype(np.int16) 160 | 161 | noise = np.zeros_like(raw_noise) 162 | 163 | if poisoned_indices_path is not None: 164 | with open(poisoned_indices_path, 'rb') as f: 165 | indices = pickle.load(f) 166 | else: 167 | indices = np.random.permutation(len(noise))[:int(len(noise)*noise_rate)] 168 | 169 | noise[indices] += raw_noise[indices] 170 | 171 | ''' restore noise (NCWH) for raw images (NHWC) ''' 172 | noise = np.transpose(noise, [0,2,3,1]) 173 | 174 | ''' add noise to images (uint8, 0~255) ''' 175 | imgs = target_set.x.astype(np.int16) + noise 176 | imgs = imgs.clip(0,255).astype(np.uint8) 177 | target_set.x = imgs 178 | 179 | target_set = data.Dataset(x=target_set.x, y=target_set.y, transform=target_set.transform, fitr=target_set.fitr) 180 | 181 | if train: 182 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 183 | else: 184 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 185 | 186 | return loader 187 | 188 | 189 | def get_clear_loader( 190 | dataset, batch_size, root='./data', train=True, 191 | noise_rate=1.0, poisoned_indices_path=None, fitr=None): 192 | 193 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 194 | return imagenet_utils.get_clear_loader( 195 | dataset, batch_size, root, train, noise_rate, poisoned_indices_path) 196 | 197 | target_set = get_dataset(dataset, root=root, train=train, fitr=fitr) 198 | data_nums = len(target_set) 199 | 200 | if poisoned_indices_path is not None: 201 | with open(poisoned_indices_path, 'rb') as f: 202 | poi_indices = pickle.load(f) 203 | indices = np.array( list( set(range(data_nums)) - set(poi_indices) ) ) 204 | 205 | else: 206 | indices = np.random.permutation(range(data_nums))[: int( data_nums * (1-noise_rate) )] 207 | 208 | ''' select clear examples ''' 209 | target_set.x = target_set.x[indices] 210 | target_set.y = np.array(target_set.y)[indices] 211 | 212 | target_set = data.Dataset(x=target_set.x, y=target_set.y, transform=target_set.transform, fitr=target_set.fitr) 213 | 214 | if train: 215 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=True, drop_last=True) 216 | else: 217 | loader = data.Loader(target_set, batch_size=batch_size, shuffle=False, drop_last=False) 218 | 219 | return loader 220 | 221 | 222 | def get_arch(arch, dataset): 223 | if dataset == 'cifar10': 224 | in_dims, out_dims = 3, 10 225 | elif dataset == 'cifar100': 226 | in_dims, out_dims = 3, 100 227 | elif dataset == 'tiny-imagenet': 228 | in_dims, out_dims = 3, 200 229 | elif dataset == 'imagenet': 230 | in_dims, out_dims = 3, 1000 231 | elif dataset == 'imagenet-mini': 232 | in_dims, out_dims = 3, 100 233 | else: 234 | raise NotImplementedError('dataset {} is not supported'.format(dataset)) 235 | 236 | if arch == 'resnet18': 237 | return models.resnet18(in_dims, out_dims) 238 | 239 | elif arch == 'resnet50': 240 | return models.resnet50(in_dims, out_dims) 241 | 242 | elif arch == 'wrn-34-10': 243 | return models.wrn34_10(in_dims, out_dims) 244 | 245 | elif arch == 'vgg11-bn': 246 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 247 | raise NotImplementedError 248 | return models.vgg11_bn(in_dims, out_dims) 249 | 250 | elif arch == 'vgg16-bn': 251 | if dataset == 'imagenet' or dataset == 'imagenet-mini': 252 | raise NotImplementedError 253 | return models.vgg16_bn(in_dims, out_dims) 254 | 255 | elif arch == 'vgg19-bn': 256 | return models.vgg19_bn(in_dims, out_dims) 257 | 258 | elif arch == 'densenet-121': 259 | return models.densenet121(num_classes=out_dims) 260 | 261 | else: 262 | raise NotImplementedError('architecture {} is not supported'.format(arch)) 263 | 264 | 265 | def get_optim(optim, params, lr=0.1, weight_decay=1e-4, momentum=0.9): 266 | if optim == 'sgd': 267 | return torch.optim.SGD(params, lr=lr, weight_decay=weight_decay, momentum=momentum) 268 | elif optim == 'adam': 269 | return torch.optim.Adam(params, lr=lr, weight_decay=weight_decay) 270 | 271 | raise NotImplementedError('optimizer {} is not supported'.format(optim)) 272 | 273 | 274 | def generic_init(args): 275 | if os.path.exists(args.save_dir) == False: 276 | os.makedirs(args.save_dir) 277 | 278 | fmt = '%(asctime)s %(name)s:%(levelname)s: %(message)s' 279 | formatter = logging.Formatter( 280 | fmt, datefmt='%Y-%m-%d %H:%M:%S') 281 | 282 | fh = logging.FileHandler( 283 | '{}/{}_log.txt'.format(args.save_dir, args.save_name), mode='w') 284 | fh.setFormatter(formatter) 285 | 286 | logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=fmt, datefmt='%Y-%m-%d %H:%M:%S') 287 | logger = logging.getLogger() 288 | logger.addHandler(fh) 289 | 290 | logger.info('Arguments') 291 | for arg in vars(args): 292 | logger.info(' {:<22} {}'.format(arg+':', getattr(args,arg)) ) 293 | logger.info('') 294 | 295 | return logger 296 | 297 | 298 | def evaluate(model, criterion, loader, cpu): 299 | acc = AverageMeter() 300 | loss = AverageMeter() 301 | 302 | model.eval() 303 | for x, y in loader: 304 | if not cpu: x, y = x.cuda(), y.cuda() 305 | with torch.no_grad(): 306 | _y = model(x) 307 | ac = (_y.argmax(dim=1) == y).sum().item() / len(x) 308 | lo = criterion(_y,y).item() 309 | acc.update(ac, len(x)) 310 | loss.update(lo, len(x)) 311 | 312 | return acc.average(), loss.average() 313 | 314 | 315 | def model_state_to_cpu(model_state): 316 | model_state_cpu = type(model_state)() 317 | for key, val in model_state.items(): 318 | model_state_cpu[key] = val.cpu() 319 | return model_state_cpu 320 | 321 | 322 | def get_model_state(model): 323 | # if isinstance(model, torch.nn.parallel.DistributedDataParallel): 324 | if isinstance(model, torch.nn.DataParallel): 325 | model_state = model_state_to_cpu(model.module.state_dict()) 326 | else: 327 | model_state = model.state_dict() 328 | return model_state 329 | -------------------------------------------------------------------------------- /models/densenet.py: -------------------------------------------------------------------------------- 1 | ''' adapted from 2 | https://github.com/pytorch/vision/blob/main/torchvision/models/densenet.py 3 | ''' 4 | 5 | import re 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torch.utils.checkpoint as cp 10 | from collections import OrderedDict 11 | from torch import Tensor 12 | from typing import Any, List, Tuple 13 | 14 | 15 | __all__ = ['DenseNet', 'densenet121', 'densenet169', 'densenet201', 'densenet161'] 16 | 17 | 18 | class _DenseLayer(nn.Module): 19 | def __init__( 20 | self, 21 | num_input_features: int, 22 | growth_rate: int, 23 | bn_size: int, 24 | drop_rate: float, 25 | memory_efficient: bool = False 26 | ) -> None: 27 | super(_DenseLayer, self).__init__() 28 | self.norm1: nn.BatchNorm2d 29 | self.add_module('norm1', nn.BatchNorm2d(num_input_features)) 30 | self.relu1: nn.ReLU 31 | self.add_module('relu1', nn.ReLU(inplace=True)) 32 | self.conv1: nn.Conv2d 33 | self.add_module('conv1', nn.Conv2d(num_input_features, bn_size * 34 | growth_rate, kernel_size=1, stride=1, 35 | bias=False)) 36 | self.norm2: nn.BatchNorm2d 37 | self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)) 38 | self.relu2: nn.ReLU 39 | self.add_module('relu2', nn.ReLU(inplace=True)) 40 | self.conv2: nn.Conv2d 41 | self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate, 42 | kernel_size=3, stride=1, padding=1, 43 | bias=False)) 44 | self.drop_rate = float(drop_rate) 45 | self.memory_efficient = memory_efficient 46 | 47 | def bn_function(self, inputs: List[Tensor]) -> Tensor: 48 | concated_features = torch.cat(inputs, 1) 49 | bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features))) # noqa: T484 50 | return bottleneck_output 51 | 52 | # todo: rewrite when torchscript supports any 53 | def any_requires_grad(self, input: List[Tensor]) -> bool: 54 | for tensor in input: 55 | if tensor.requires_grad: 56 | return True 57 | return False 58 | 59 | @torch.jit.unused # noqa: T484 60 | def call_checkpoint_bottleneck(self, input: List[Tensor]) -> Tensor: 61 | def closure(*inputs): 62 | return self.bn_function(inputs) 63 | 64 | return cp.checkpoint(closure, *input) 65 | 66 | @torch.jit._overload_method # noqa: F811 67 | def forward(self, input: List[Tensor]) -> Tensor: 68 | pass 69 | 70 | @torch.jit._overload_method # noqa: F811 71 | def forward(self, input: Tensor) -> Tensor: 72 | pass 73 | 74 | # torchscript does not yet support *args, so we overload method 75 | # allowing it to take either a List[Tensor] or single Tensor 76 | def forward(self, input: Tensor) -> Tensor: # noqa: F811 77 | if isinstance(input, Tensor): 78 | prev_features = [input] 79 | else: 80 | prev_features = input 81 | 82 | if self.memory_efficient and self.any_requires_grad(prev_features): 83 | if torch.jit.is_scripting(): 84 | raise Exception("Memory Efficient not supported in JIT") 85 | 86 | bottleneck_output = self.call_checkpoint_bottleneck(prev_features) 87 | else: 88 | bottleneck_output = self.bn_function(prev_features) 89 | 90 | new_features = self.conv2(self.relu2(self.norm2(bottleneck_output))) 91 | if self.drop_rate > 0: 92 | new_features = F.dropout(new_features, p=self.drop_rate, 93 | training=self.training) 94 | return new_features 95 | 96 | 97 | class _DenseBlock(nn.ModuleDict): 98 | _version = 2 99 | 100 | def __init__( 101 | self, 102 | num_layers: int, 103 | num_input_features: int, 104 | bn_size: int, 105 | growth_rate: int, 106 | drop_rate: float, 107 | memory_efficient: bool = False 108 | ) -> None: 109 | super(_DenseBlock, self).__init__() 110 | for i in range(num_layers): 111 | layer = _DenseLayer( 112 | num_input_features + i * growth_rate, 113 | growth_rate=growth_rate, 114 | bn_size=bn_size, 115 | drop_rate=drop_rate, 116 | memory_efficient=memory_efficient, 117 | ) 118 | self.add_module('denselayer%d' % (i + 1), layer) 119 | 120 | def forward(self, init_features: Tensor) -> Tensor: 121 | features = [init_features] 122 | for name, layer in self.items(): 123 | new_features = layer(features) 124 | features.append(new_features) 125 | return torch.cat(features, 1) 126 | 127 | 128 | class _Transition(nn.Sequential): 129 | def __init__(self, num_input_features: int, num_output_features: int) -> None: 130 | super(_Transition, self).__init__() 131 | self.add_module('norm', nn.BatchNorm2d(num_input_features)) 132 | self.add_module('relu', nn.ReLU(inplace=True)) 133 | self.add_module('conv', nn.Conv2d(num_input_features, num_output_features, 134 | kernel_size=1, stride=1, bias=False)) 135 | self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2)) 136 | 137 | 138 | class DenseNet(nn.Module): 139 | r"""Densenet-BC model class, based on 140 | `"Densely Connected Convolutional Networks" `_. 141 | 142 | Args: 143 | growth_rate (int) - how many filters to add each layer (`k` in paper) 144 | block_config (list of 4 ints) - how many layers in each pooling block 145 | num_init_features (int) - the number of filters to learn in the first convolution layer 146 | bn_size (int) - multiplicative factor for number of bottle neck layers 147 | (i.e. bn_size * k features in the bottleneck layer) 148 | drop_rate (float) - dropout rate after each dense layer 149 | num_classes (int) - number of classification classes 150 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 151 | but slower. Default: *False*. See `"paper" `_. 152 | """ 153 | 154 | def __init__( 155 | self, 156 | growth_rate: int = 32, 157 | block_config: Tuple[int, int, int, int] = (6, 12, 24, 16), 158 | num_init_features: int = 64, 159 | bn_size: int = 4, 160 | drop_rate: float = 0, 161 | num_classes: int = 1000, 162 | memory_efficient: bool = False 163 | ) -> None: 164 | 165 | super(DenseNet, self).__init__() 166 | 167 | # First convolution 168 | self.features = nn.Sequential(OrderedDict([ 169 | ('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, 170 | padding=3, bias=False)), 171 | ('norm0', nn.BatchNorm2d(num_init_features)), 172 | ('relu0', nn.ReLU(inplace=True)), 173 | ('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)), 174 | ])) 175 | 176 | # Each denseblock 177 | num_features = num_init_features 178 | for i, num_layers in enumerate(block_config): 179 | block = _DenseBlock( 180 | num_layers=num_layers, 181 | num_input_features=num_features, 182 | bn_size=bn_size, 183 | growth_rate=growth_rate, 184 | drop_rate=drop_rate, 185 | memory_efficient=memory_efficient 186 | ) 187 | self.features.add_module('denseblock%d' % (i + 1), block) 188 | num_features = num_features + num_layers * growth_rate 189 | if i != len(block_config) - 1: 190 | trans = _Transition(num_input_features=num_features, 191 | num_output_features=num_features // 2) 192 | self.features.add_module('transition%d' % (i + 1), trans) 193 | num_features = num_features // 2 194 | 195 | # Final batch norm 196 | self.features.add_module('norm5', nn.BatchNorm2d(num_features)) 197 | 198 | # Linear layer 199 | self.classifier = nn.Linear(num_features, num_classes) 200 | 201 | # Official init from torch repo. 202 | for m in self.modules(): 203 | if isinstance(m, nn.Conv2d): 204 | nn.init.kaiming_normal_(m.weight) 205 | elif isinstance(m, nn.BatchNorm2d): 206 | nn.init.constant_(m.weight, 1) 207 | nn.init.constant_(m.bias, 0) 208 | elif isinstance(m, nn.Linear): 209 | nn.init.constant_(m.bias, 0) 210 | 211 | def forward(self, x: Tensor) -> Tensor: 212 | features = self.features(x) 213 | out = F.relu(features, inplace=True) 214 | out = F.adaptive_avg_pool2d(out, (1, 1)) 215 | out = torch.flatten(out, 1) 216 | out = self.classifier(out) 217 | return out 218 | 219 | 220 | def _densenet( 221 | arch: str, 222 | growth_rate: int, 223 | block_config: Tuple[int, int, int, int], 224 | num_init_features: int, 225 | pretrained: bool, 226 | progress: bool, 227 | **kwargs: Any 228 | ) -> DenseNet: 229 | model = DenseNet(growth_rate, block_config, num_init_features, **kwargs) 230 | if pretrained: 231 | raise NotImplementedError 232 | return model 233 | 234 | 235 | def densenet121(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DenseNet: 236 | r"""Densenet-121 model from 237 | `"Densely Connected Convolutional Networks" `_. 238 | The required minimum input size of the model is 29x29. 239 | 240 | Args: 241 | pretrained (bool): If True, returns a model pre-trained on ImageNet 242 | progress (bool): If True, displays a progress bar of the download to stderr 243 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 244 | but slower. Default: *False*. See `"paper" `_. 245 | """ 246 | return _densenet('densenet121', 32, (6, 12, 24, 16), 64, pretrained, progress, 247 | **kwargs) 248 | 249 | 250 | def densenet161(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DenseNet: 251 | r"""Densenet-161 model from 252 | `"Densely Connected Convolutional Networks" `_. 253 | The required minimum input size of the model is 29x29. 254 | 255 | Args: 256 | pretrained (bool): If True, returns a model pre-trained on ImageNet 257 | progress (bool): If True, displays a progress bar of the download to stderr 258 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 259 | but slower. Default: *False*. See `"paper" `_. 260 | """ 261 | return _densenet('densenet161', 48, (6, 12, 36, 24), 96, pretrained, progress, 262 | **kwargs) 263 | 264 | 265 | def densenet169(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DenseNet: 266 | r"""Densenet-169 model from 267 | `"Densely Connected Convolutional Networks" `_. 268 | The required minimum input size of the model is 29x29. 269 | 270 | Args: 271 | pretrained (bool): If True, returns a model pre-trained on ImageNet 272 | progress (bool): If True, displays a progress bar of the download to stderr 273 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 274 | but slower. Default: *False*. See `"paper" `_. 275 | """ 276 | return _densenet('densenet169', 32, (6, 12, 32, 32), 64, pretrained, progress, 277 | **kwargs) 278 | 279 | 280 | def densenet201(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DenseNet: 281 | r"""Densenet-201 model from 282 | `"Densely Connected Convolutional Networks" `_. 283 | The required minimum input size of the model is 29x29. 284 | 285 | Args: 286 | pretrained (bool): If True, returns a model pre-trained on ImageNet 287 | progress (bool): If True, displays a progress bar of the download to stderr 288 | memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient, 289 | but slower. Default: *False*. See `"paper" `_. 290 | """ 291 | return _densenet('densenet201', 32, (6, 12, 48, 32), 64, pretrained, progress, 292 | **kwargs) 293 | --------------------------------------------------------------------------------