├── .DS_Store ├── .gitignore ├── Attack_BadNets.py ├── Attack_Blended.py ├── Defense_ShrinkPad.py ├── LICENSE ├── README.md ├── core ├── .DS_Store ├── __init__.py ├── attacks │ ├── AdaptivePatch.py │ ├── BATT.py │ ├── BadNets.py │ ├── Blended.py │ ├── Blind.py │ ├── IAD.py │ ├── ISSBA.py │ ├── LIRA.py │ ├── LabelConsistent.py │ ├── PhysicalBA.py │ ├── Refool.py │ ├── SleeperAgent.py │ ├── TUAP.py │ ├── WaNet.py │ ├── __init__.py │ └── base.py ├── defenses │ ├── ABL.py │ ├── AutoEncoderDefense.py │ ├── CutMix.py │ ├── FineTuning.py │ ├── IBD_PSC.py │ ├── MCR.py │ ├── NAD.py │ ├── Pruning.py │ ├── REFINE.py │ ├── SCALE_UP.py │ ├── ShrinkPad.py │ ├── Spectral.py │ ├── __init__.py │ └── base.py ├── models │ ├── __init__.py │ ├── autoencoder.py │ ├── baseline_MNIST_network.py │ ├── curves.py │ ├── resnet.py │ ├── resnet_curve.py │ ├── unet.py │ ├── vgg.py │ └── vgg_curve.py └── utils │ ├── __init__.py │ ├── accuracy.py │ ├── any2tensor.py │ ├── compute_metric.py │ ├── log.py │ ├── supconloss.py │ ├── test.py │ └── torchattacks │ ├── README.md │ ├── __init__.py │ ├── attack.py │ └── attacks │ ├── __init__.py │ └── pgd.py ├── requirements.txt └── tests ├── .DS_Store ├── test_ABL.py ├── test_AdaptivePatch.py ├── test_AutoEncoder.py ├── test_BATT.py ├── test_BadNets.py ├── test_BadNets_with_multi_gpus.py ├── test_Blended.py ├── test_Blind.py ├── test_CutMix.py ├── test_FineTuning.py ├── test_IAD.py ├── test_IBD-PSC.py ├── test_ISSBA.py ├── test_LIRA.py ├── test_LabelConsistent.py ├── test_MCR.py ├── test_NAD.py ├── test_PhysicalBA.py ├── test_Pruning.py ├── test_REFINE.py ├── test_SCALE_UP.py ├── test_ShrinkPad.py ├── test_SleeperAgent.py ├── test_Spectral.py ├── test_TUAP.py ├── test_WaNet.py ├── test_cifar10.py ├── test_dataset_folder.py ├── test_refool.py ├── test_refool_vis.ipynb ├── train_AutoEncoder.py ├── train_Benign.py └── train_REFINE.py /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUYimingLi/BackdoorBox/945ea2125096dd19720be1489d0d92239c0541fd/.DS_Store -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # data 7 | data/ 8 | adv_dataset_old/ 9 | 10 | # experiments 11 | experiments/ 12 | 13 | # personal test code 14 | *data*/ 15 | *experiments*/ 16 | tests/model/ 17 | mytest/ 18 | ./test.sh 19 | ./test.py 20 | ./test2.py 21 | ./result 22 | ./experiments -------------------------------------------------------------------------------- /Attack_BadNets.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the example code of benign training and poisoned training on torchvision.datasets.DatasetFolder. 3 | Dataset is CIFAR-10. 4 | Attack method is BadNets. 5 | ''' 6 | 7 | 8 | import os 9 | 10 | import cv2 11 | import torch 12 | import torch.nn as nn 13 | from torch.utils.data import Dataset 14 | import torchvision 15 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 16 | 17 | import core 18 | 19 | 20 | 21 | dataset = torchvision.datasets.DatasetFolder 22 | 23 | # image file -> cv.imread -> numpy.ndarray (H x W x C) -> ToTensor -> torch.Tensor (C x H x W) -> RandomHorizontalFlip -> torch.Tensor -> network input 24 | transform_train = Compose([ 25 | ToTensor(), 26 | RandomHorizontalFlip() 27 | ]) 28 | 29 | trainset = dataset( 30 | root='./data/cifar10/train', 31 | loader=cv2.imread, 32 | extensions=('png',), 33 | transform=transform_train, 34 | target_transform=None, 35 | is_valid_file=None) 36 | 37 | transform_test = Compose([ 38 | ToTensor() 39 | ]) 40 | testset = dataset( 41 | root='./data/cifar10/test', 42 | loader=cv2.imread, 43 | extensions=('png',), 44 | transform=transform_test, 45 | target_transform=None, 46 | is_valid_file=None) 47 | 48 | index = 44 49 | 50 | x, y = trainset[index] 51 | print(y) 52 | for a in x[0]: 53 | for b in a: 54 | print("%-4.2f" % float(b), end=' ') 55 | print() 56 | 57 | 58 | pattern = torch.zeros((1, 32, 32), dtype=torch.uint8) 59 | pattern[0, -3:, -3:] = 255 60 | weight = torch.zeros((1, 32, 32), dtype=torch.float32) 61 | weight[0, -3:, -3:] = 1.0 62 | 63 | badnets = core.BadNets( 64 | train_dataset=trainset, 65 | test_dataset=testset, 66 | model=core.models.ResNet(18), 67 | # model=core.models.BaselineMNISTNetwork(), 68 | loss=nn.CrossEntropyLoss(), 69 | y_target=0, 70 | poisoned_rate=0.1, 71 | pattern=pattern, 72 | weight=weight, 73 | # poisoned_transform_index=0, 74 | # poisoned_target_transform_index=0, 75 | schedule=None, 76 | seed=666 77 | ) 78 | 79 | poisoned_train_dataset, poisoned_test_dataset = badnets.get_poisoned_dataset() 80 | 81 | x, y = poisoned_train_dataset[index] 82 | print(y) 83 | for a in x[0]: 84 | for b in a: 85 | print("%-4.2f" % float(b), end=' ') 86 | print() 87 | 88 | x, y = poisoned_test_dataset[index] 89 | print(y) 90 | for a in x[0]: 91 | for b in a: 92 | print("%-4.2f" % float(b), end=' ') 93 | print() 94 | 95 | # train benign model 96 | schedule = { 97 | 'device': 'GPU', 98 | 'CUDA_VISIBLE_DEVICES': '0', 99 | 'GPU_num': 1, 100 | 101 | 'benign_training': True, 102 | 'batch_size': 128, 103 | 'num_workers': 16, 104 | 105 | 'lr': 0.1, 106 | 'momentum': 0.9, 107 | 'weight_decay': 5e-4, 108 | 'gamma': 0.1, 109 | 'schedule': [150, 180], 110 | 111 | 'epochs': 200, 112 | 113 | 'log_iteration_interval': 100, 114 | 'test_epoch_interval': 10, 115 | 'save_epoch_interval': 10, 116 | 117 | 'save_dir': 'experiments', 118 | 'experiment_name': 'train_benign_DatasetFolder-CIFAR10' 119 | } 120 | 121 | badnets.train(schedule) 122 | 123 | # train attacked model 124 | schedule = { 125 | 'device': 'GPU', 126 | 'CUDA_VISIBLE_DEVICES': '0', 127 | 'GPU_num': 1, 128 | 129 | 'benign_training': False, 130 | 'batch_size': 128, 131 | 'num_workers': 16, 132 | 133 | 'lr': 0.1, 134 | 'momentum': 0.9, 135 | 'weight_decay': 5e-4, 136 | 'gamma': 0.1, 137 | 'schedule': [150, 180], 138 | 139 | 'epochs': 200, 140 | 141 | 'log_iteration_interval': 100, 142 | 'test_epoch_interval': 10, 143 | 'save_epoch_interval': 10, 144 | 145 | 'save_dir': 'experiments', 146 | 'experiment_name': 'train_poisoned_DatasetFolder-CIFAR10' 147 | } 148 | 149 | badnets.train(schedule) -------------------------------------------------------------------------------- /Attack_Blended.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the example code of benign training and poisoned training on torchvision.datasets.DatasetFolder. 3 | Dataset is ImageNet50. 4 | Attack method is Blended. 5 | ''' 6 | 7 | import os 8 | 9 | import torch 10 | import torch.nn as nn 11 | from torch.utils.data import Dataset 12 | import torchvision 13 | from torchvision.transforms import Compose, ToTensor, RandomResizedCrop, RandomHorizontalFlip, Normalize 14 | import core 15 | 16 | global_seed = 666 17 | deterministic = True 18 | torch.manual_seed(global_seed) 19 | 20 | # Change dataset to ImageNet50 21 | dataset = torchvision.datasets.ImageFolder 22 | 23 | transform_train = Compose([ 24 | RandomResizedCrop(224), 25 | RandomHorizontalFlip(), 26 | ToTensor(), 27 | Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), 28 | ]) 29 | trainset = dataset(root='./data/ImageNet50/train', transform=transform_train) 30 | 31 | transform_test = Compose([ 32 | RandomResizedCrop(224), 33 | ToTensor(), 34 | Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), 35 | ]) 36 | testset = dataset(root='./data/ImageNet50/val', transform=transform_test) 37 | 38 | # Update pattern and weight for ImageNet50 (224x224 images) 39 | pattern = torch.zeros((3, 224, 224), dtype=torch.uint8) 40 | pattern[:, -3:, -3:] = 255 41 | weight = torch.zeros((3, 224, 224), dtype=torch.float32) 42 | weight[:, -3:, -3:] = 0.2 43 | 44 | blended = core.Blended( 45 | train_dataset=trainset, 46 | test_dataset=testset, 47 | model=core.models.ResNet(18, num_classes=50), # Use ResNet for ImageNet50 48 | loss=nn.CrossEntropyLoss(), 49 | pattern=pattern, 50 | weight=weight, 51 | y_target=1, 52 | poisoned_rate=0.1, # 设置投毒率为10% 53 | seed=global_seed, 54 | deterministic=deterministic 55 | ) 56 | 57 | poisoned_train_dataset, poisoned_test_dataset = blended.get_poisoned_dataset() 58 | 59 | index = 0 # 确保在使用之前定义 index 60 | # Show an Example of Poisoned Training Samples 61 | x, y = poisoned_train_dataset[index] 62 | print(y) 63 | for a in x[0]: 64 | for b in a: 65 | print("%-4.2f" % float(b), end=' ') 66 | print() 67 | 68 | 69 | # Show an Example of Poisoned Testing Samples 70 | x, y = poisoned_test_dataset[index] 71 | print(y) 72 | for a in x[0]: 73 | for b in a: 74 | print("%-4.2f" % float(b), end=' ') 75 | print() 76 | 77 | 78 | # Train Benign Model 79 | schedule = { 80 | 'device': 'GPU', 81 | 'CUDA_VISIBLE_DEVICES': '3', 82 | 'GPU_num': 1, 83 | 84 | 'benign_training': True, 85 | 'batch_size': 128, 86 | 'num_workers': 4, 87 | 88 | 'lr': 0.1, 89 | 'momentum': 0.9, 90 | 'weight_decay': 5e-4, 91 | 'gamma': 0.1, 92 | 'schedule': [30, 60, 90], 93 | 94 | 'epochs': 100, 95 | 96 | 'log_iteration_interval': 100, 97 | 'test_epoch_interval': 10, 98 | 'save_epoch_interval': 10, 99 | 100 | 'save_dir': 'experiments', 101 | 'experiment_name': 'train_benign_ImageNet50_Blended' 102 | } 103 | 104 | blended.train(schedule) 105 | benign_model = blended.get_model() 106 | torch.save(benign_model.state_dict(), 'benign_model_resnet18.pt') 107 | 108 | # Test Benign Model 109 | test_schedule = { 110 | 'device': 'GPU', 111 | 'CUDA_VISIBLE_DEVICES': '3', 112 | 'GPU_num': 1, 113 | 114 | 'batch_size': 128, 115 | 'num_workers': 4, 116 | 117 | 'save_dir': 'experiments', 118 | 'experiment_name': 'test_benign_ImageNet50_Blended' 119 | } 120 | blended.test(test_schedule) 121 | torch.save(benign_model.state_dict(), 'benign_model_resnet18_tested.pt') 122 | 123 | 124 | blended.model = core.models.ResNet(18, num_classes=50) 125 | # Train Infected Model 126 | schedule = { 127 | 'device': 'GPU', 128 | 'CUDA_VISIBLE_DEVICES': '3', 129 | 'GPU_num': 1, 130 | 131 | 'benign_training': False, 132 | 'batch_size': 128, 133 | 'num_workers': 4, 134 | 135 | 'lr': 0.1, 136 | 'momentum': 0.9, 137 | 'weight_decay': 5e-4, 138 | 'gamma': 0.1, 139 | 'schedule': [30, 60, 90], 140 | 141 | 'epochs': 100, 142 | 143 | 'log_iteration_interval': 100, 144 | 'test_epoch_interval': 10, 145 | 'save_epoch_interval': 10, 146 | 147 | 'save_dir': 'experiments', 148 | 'experiment_name': 'train_poisoned_ImageNet50_Blended' 149 | } 150 | 151 | blended.train(schedule) 152 | infected_model = blended.get_model() 153 | torch.save(infected_model.state_dict(), 'infected_model_resnet18.pt') 154 | 155 | # Test Infected Model 156 | test_schedule = { 157 | 'device': 'GPU', 158 | 'CUDA_VISIBLE_DEVICES': '3', 159 | 'GPU_num': 1, 160 | 161 | 'batch_size': 128, 162 | 'num_workers': 4, 163 | 164 | 'save_dir': 'experiments', 165 | 'experiment_name': 'test_poisoned_ImageNet50_Blended' 166 | } 167 | blended.test(test_schedule) 168 | torch.save(infected_model.state_dict(), 'infected_model_resnet18_tested.pt') -------------------------------------------------------------------------------- /Defense_ShrinkPad.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the example code of defending the BadNets attack. 3 | Dataset is CIFAR-10. 4 | Defense method is ShrinkPad. 5 | ''' 6 | 7 | import os 8 | 9 | import cv2 10 | import torch 11 | import torch.nn as nn 12 | from torch.utils.data import Dataset 13 | import torchvision.models as models 14 | import torchvision 15 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip, Lambda 16 | 17 | import core 18 | 19 | shrinkpad = core.ShrinkPad(size_map=32, pad=4) 20 | 21 | dataset = torchvision.datasets.DatasetFolder 22 | 23 | pattern = torch.zeros((1, 32, 32), dtype=torch.uint8) 24 | pattern[0, -3:, -3:] = 255 25 | 26 | # The targeting models that have been poisoned and trained by the BadNet attack, the specific watermark is a small 3x3 square, which you can see in example.py. 27 | 28 | transform_test = Compose([ 29 | ToTensor() , 30 | Lambda(lambda img: img + pattern) , 31 | Lambda(lambda img: shrinkpad.preprocess(img)) 32 | ]) 33 | 34 | testset = dataset( 35 | root='./data/cifar10/test', 36 | loader=cv2.imread, 37 | extensions=('png',), 38 | transform=transform_test, 39 | target_transform=None, 40 | is_valid_file=None) 41 | 42 | schedule = { 43 | 'test_model': './experiments/train_poisoned_DatasetFolder-CIFAR10_2025-02-24_18:20:13/ckpt_epoch_50.pth', 44 | 'save_dir': './experiments', 45 | 'CUDA_VISIBLE_DEVICES': '0', 46 | 'GPU_num': 1, 47 | 'experiment_name': 'CIFAR10_test', 48 | 'device': 'GPU', 49 | 'metric': 'ASR_NoTarget', 50 | 'y_target': 0, 51 | 'batch_size': 64, 52 | 'num_workers': 4, 53 | } 54 | model=core.models.ResNet(18) 55 | 56 | predictions = shrinkpad.test(model, testset, schedule=schedule, size_map=32, pad=4) 57 | 58 | -------------------------------------------------------------------------------- /core/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUYimingLi/BackdoorBox/945ea2125096dd19720be1489d0d92239c0541fd/core/.DS_Store -------------------------------------------------------------------------------- /core/__init__.py: -------------------------------------------------------------------------------- 1 | from .attacks import * 2 | from .defenses import * 3 | from .models import * 4 | -------------------------------------------------------------------------------- /core/attacks/__init__.py: -------------------------------------------------------------------------------- 1 | from ast import Import 2 | from .BadNets import BadNets 3 | from .Blended import Blended 4 | from .LabelConsistent import LabelConsistent 5 | from .Refool import Refool 6 | from .WaNet import WaNet 7 | from .Blind import Blind 8 | from .IAD import IAD 9 | from .LIRA import LIRA 10 | from .PhysicalBA import PhysicalBA 11 | from .ISSBA import ISSBA 12 | from .TUAP import TUAP 13 | from .SleeperAgent import SleeperAgent 14 | from .BATT import BATT 15 | from .AdaptivePatch import AdaptivePatch 16 | 17 | __all__ = [ 18 | 'BadNets', 'Blended','Refool', 'WaNet', 'LabelConsistent', 'Blind', 'IAD', 'LIRA', 'PhysicalBA', 'ISSBA','TUAP', 'SleeperAgent','BATT', 'AdaptivePatch' 19 | ] 20 | -------------------------------------------------------------------------------- /core/defenses/FineTuning.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the implement of fine-tuning proposed in [1]. 3 | [1] Fine-Pruning: Defending Against Backdooring Attacks on Deep Neural Networks. RAID, 2018. 4 | ''' 5 | 6 | 7 | import os 8 | 9 | import torch 10 | import torch.nn as nn 11 | import numpy as np 12 | 13 | from .base import Base 14 | from ..utils import test 15 | from torch.utils.data import DataLoader 16 | import random 17 | import time 18 | import os.path as osp 19 | from ..utils import Log 20 | 21 | 22 | def _seed_worker(worker_id): 23 | worker_seed = torch.initial_seed() % 2 ** 32 24 | np.random.seed(worker_seed) 25 | random.seed(worker_seed) 26 | 27 | 28 | def adjust_learning_rate(lr, optimizer, epoch): 29 | if epoch in [20]: 30 | lr *= 0.1 31 | for param_group in optimizer.param_groups: 32 | param_group['lr'] = lr 33 | 34 | 35 | 36 | class FineTuning(Base): 37 | """FineTuning process. 38 | Args: 39 | train_dataset (types in support_list): Benign training dataset. 40 | test_dataset (types in support_list): Benign testing dataset. 41 | model (torch.nn.Module): Network. 42 | layer(list): The layers to fintune 43 | loss (torch.nn.Module): Loss. 44 | schedule (dict): Training or testing schedule. Default: None. 45 | seed (int): Global seed for random numbers. Default: 0. 46 | deterministic (bool): Sets whether PyTorch operations must use "deterministic" algorithms. 47 | That is, algorithms which, given the same input, and when run on the same software and hardware, 48 | always produce the same output. When enabled, operations will use deterministic algorithms when available, 49 | and if only nondeterministic algorithms are available they will throw a RuntimeError when called. Default: False. 50 | """ 51 | def __init__(self, 52 | train_dataset=None, 53 | test_dataset=None, 54 | model=None, 55 | layer=None, 56 | loss=None, 57 | schedule=None, 58 | seed=0, 59 | deterministic=False): 60 | 61 | super(FineTuning, self).__init__(seed=seed, deterministic=deterministic) 62 | 63 | self.train_dataset=train_dataset 64 | self.test_dataset=test_dataset 65 | self.model=model 66 | self.layer=layer 67 | self.loss=loss 68 | self.schedule=schedule 69 | 70 | 71 | def frozen(self): 72 | """Frozen the layers which don't need to fine tuning. 73 | """ 74 | if self.layer==None or self.layer[0]=="full layers": 75 | return 76 | else: 77 | for name, child in self.model.named_children(): 78 | if not name in self.layer: 79 | for param in child.parameters(): 80 | param.requires_grad = False 81 | 82 | def repair(self,schedule=None): 83 | """Finetuning. 84 | Args: 85 | schedule (dict): Schedule for testing. 86 | """ 87 | self.frozen() 88 | print("--------fine tuning-------") 89 | if schedule==None: 90 | raise AttributeError("Schedule is None, please check your schedule setting.") 91 | current_schedule =schedule 92 | 93 | # Use GPU 94 | if 'device' in current_schedule and current_schedule['device'] == 'GPU': 95 | if 'CUDA_VISIBLE_DEVICES' in current_schedule: 96 | os.environ['CUDA_VISIBLE_DEVICES'] = current_schedule['CUDA_VISIBLE_DEVICES'] 97 | 98 | assert torch.cuda.device_count() > 0, 'This machine has no cuda devices!' 99 | assert current_schedule['GPU_num'] > 0, 'GPU_num should be a positive integer' 100 | print( 101 | f"This machine has {torch.cuda.device_count()} cuda devices, and use {current_schedule['GPU_num']} of them to train.") 102 | 103 | if current_schedule['GPU_num'] == 1: 104 | device = torch.device("cuda:0") 105 | else: 106 | gpus = list(range(current_schedule['GPU_num'])) 107 | self.model = nn.DataParallel(self.model.cuda(), device_ids=gpus, output_device=gpus[0]) 108 | # TODO: DDP training 109 | pass 110 | # Use CPU 111 | else: 112 | device = torch.device("cpu") 113 | 114 | if self.train_dataset==None: 115 | raise AttributeError("Train set is None, please check your setting.") 116 | train_loader = DataLoader( 117 | self.train_dataset, 118 | batch_size=current_schedule['batch_size'], 119 | shuffle=True, 120 | num_workers=current_schedule['num_workers'], 121 | drop_last=True, 122 | pin_memory=True, 123 | worker_init_fn=_seed_worker 124 | ) 125 | 126 | model = self.model.to(device) 127 | model.train() 128 | 129 | optimizer = torch.optim.SGD(model.parameters(), lr=current_schedule['lr'], 130 | momentum=current_schedule['momentum'], 131 | weight_decay=current_schedule['weight_decay']) 132 | work_dir = osp.join(current_schedule['save_dir'], current_schedule['experiment_name'] + '_' + time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())) 133 | os.makedirs(work_dir, exist_ok=True) 134 | log = Log(osp.join(work_dir, 'log.txt')) 135 | 136 | # log and output: 137 | # 1. ouput loss and time 138 | # 2. test and output statistics 139 | # 3. save checkpoint 140 | 141 | iteration = 0 142 | last_time = time.time() 143 | 144 | for i in range(current_schedule['epochs']): 145 | adjust_learning_rate(current_schedule['lr'],optimizer, i) 146 | for batch_id, batch in enumerate(train_loader): 147 | batch_img = batch[0] 148 | batch_label = batch[1] 149 | batch_img = batch_img.to(device) 150 | batch_label = batch_label.to(device) 151 | optimizer.zero_grad() 152 | predict_digits = self.model(batch_img) 153 | loss = self.loss(predict_digits, batch_label) 154 | loss.backward() 155 | optimizer.step() 156 | 157 | iteration += 1 158 | if iteration % current_schedule['log_iteration_interval'] == 0: 159 | msg = time.strftime("[%Y-%m-%d_%H:%M:%S] ", 160 | time.localtime()) + f"Epoch:{i + 1}/{current_schedule['epochs']}, iteration:{batch_id + 1}/{len(self.train_dataset) // current_schedule['batch_size']}, lr: {current_schedule['lr']}, loss: {float(loss)}, time: {time.time() - last_time}\n" 161 | last_time = time.time() 162 | log(msg) 163 | 164 | if (i + 1) % self.current_schedule['save_epoch_interval'] == 0: 165 | self.model.eval() 166 | self.model = self.model.cpu() 167 | ckpt_model_filename = "ckpt_epoch_" + str(i+1) + ".pth" 168 | ckpt_model_path = os.path.join(work_dir, ckpt_model_filename) 169 | torch.save(self.model.state_dict(), ckpt_model_path) 170 | self.model = self.model.to(device) 171 | self.model.train() 172 | 173 | def test(self, schedule=None): 174 | """Test the finetuning model. 175 | Args: 176 | schedule (dict): Schedule for testing. 177 | """ 178 | if schedule==None: 179 | raise AttributeError("Schedule is None, please check your schedule setting.") 180 | if self.test_dataset==None: 181 | raise AttributeError("Test set is None, please check your setting.") 182 | test(self.model,self.test_dataset,schedule) 183 | 184 | def get_model(self): 185 | return self.model -------------------------------------------------------------------------------- /core/defenses/IBD_PSC.py: -------------------------------------------------------------------------------- 1 | # This is the test code of IBD-PSC defense. 2 | # IBD-PSC: Input-level Backdoor Detection via Parameter-oriented Scaling Consistency [ICML, 2024] (https://arxiv.org/abs/2405.09786) 3 | 4 | import os 5 | import pdb 6 | import torch 7 | from torchvision import transforms 8 | from sklearn import metrics 9 | from tqdm import tqdm 10 | import copy 11 | import numpy as np 12 | import torch.nn.functional as F 13 | from PIL import Image 14 | import torchvision 15 | import matplotlib.pyplot as plt 16 | from collections import Counter 17 | from torch.utils.data import Subset 18 | import torch.nn as nn 19 | import matplotlib.pyplot as plt 20 | import numpy as np 21 | 22 | 23 | from ..utils import test 24 | 25 | from .base import Base 26 | 27 | 28 | class IBD_PSC(Base): 29 | """Identify and filter malicious testing samples (IBD-PSC). 30 | 31 | Args: 32 | model (nn.Module): The original backdoored model. 33 | n (int): The hyper-parameter for the number of parameter-amplified versions of the original backdoored model by scaling up of its different BN layers. 34 | xi (float): The hyper-parameter for the error rate. 35 | T (float): The hyper-parameter for defender-specified threshold T. If PSC(x) > T , we deem it as a backdoor sample. 36 | scale (float): The hyper-parameter for amplyfying the parameters of selected BN layers. 37 | seed (int): Global seed for random numbers. Default: 0. 38 | deterministic (bool): Sets whether PyTorch operations must use "deterministic" algorithms. 39 | 40 | 41 | """ 42 | def __init__(self, model, n=5, xi=0.6, T = 0.9, scale=1.5, valset=None, seed=666, deterministic=False): 43 | super(IBD_PSC, self).__init__(seed, deterministic) 44 | self.model = model 45 | self.model.cuda() 46 | self.model.eval() 47 | self.n = n 48 | self.xi = xi 49 | self.T = T 50 | self.scale = scale 51 | self.valset = valset 52 | 53 | layer_num = self.count_BN_layers() 54 | sorted_indices = list(range(layer_num)) 55 | sorted_indices = list(reversed(sorted_indices)) 56 | self.sorted_indices = sorted_indices 57 | self.start_index = self.prob_start(self.scale, self.sorted_indices, valset=self.valset) 58 | 59 | 60 | def count_BN_layers(self): 61 | layer_num = 0 62 | for (name1, module1) in self.model.named_modules(): 63 | if isinstance(module1, torch.nn.BatchNorm2d): 64 | # if isinstance(module1, torch.nn.Conv2d): 65 | layer_num += 1 66 | return layer_num 67 | 68 | # test accuracy on the dataset 69 | def test_acc(self, dataset, schedule): 70 | """Test repaired curve model on dataset 71 | 72 | Args: 73 | dataset (types in support_list): Dataset. 74 | schedule (dict): Schedule for testing. 75 | """ 76 | model = self.model 77 | test(model, dataset, schedule) 78 | 79 | def scale_var_index(self, index_bn, scale=1.5): 80 | copy_model = copy.deepcopy(self.model) 81 | index = -1 82 | for (name1, module1) in copy_model.named_modules(): 83 | if isinstance(module1, torch.nn.BatchNorm2d): 84 | index += 1 85 | if index in index_bn: 86 | module1.weight.data *= scale 87 | module1.bias.data *= scale 88 | return copy_model 89 | def prob_start(self, scale, sorted_indices, valset): 90 | val_loader = torch.utils.data.DataLoader(valset, batch_size=128, shuffle=False) 91 | layer_num = len(sorted_indices) 92 | # layer_index: k 93 | for layer_index in range(1, layer_num): 94 | layers = sorted_indices[:layer_index] 95 | # print(layers) 96 | smodel = self.scale_var_index(layers, scale=scale) 97 | smodel.cuda() 98 | smodel.eval() 99 | 100 | total_num = 0 101 | clean_wrong = 0 102 | with torch.no_grad(): 103 | for idx, batch in enumerate(val_loader): 104 | clean_img = batch[0] 105 | labels = batch[1] 106 | clean_img = clean_img.cuda() # batch * channels * hight * width 107 | # labels = labels.cuda() # batch 108 | clean_logits = smodel(clean_img).detach().cpu() 109 | clean_pred = torch.argmax(clean_logits, dim=1)# model prediction 110 | 111 | clean_wrong += torch.sum(labels != clean_pred) 112 | total_num += labels.shape[0] 113 | wrong_acc = clean_wrong / total_num 114 | # print(f'wrong_acc: {wrong_acc}') 115 | if wrong_acc > self.xi: 116 | return layer_index 117 | 118 | def _test(self, dataset): 119 | data_loader = torch.utils.data.DataLoader(dataset, batch_size=128, shuffle=False) 120 | self.model.eval() 121 | total_num = 0 122 | all_psc_score = [] 123 | pred_correct_mask = [] 124 | 125 | with torch.no_grad(): 126 | for idx, batch in enumerate(data_loader): 127 | imgs = batch[0] 128 | labels = batch[1] 129 | total_num += labels.shape[0] 130 | imgs = imgs.cuda() # batch * channels * hight * width 131 | labels = labels.cuda() # batch 132 | original_pred = torch.argmax(self.model(imgs), dim=1) # model prediction 133 | mask = torch.eq(labels, original_pred) # only look at those samples that successfully attack the DNN 134 | pred_correct_mask.append(mask) 135 | 136 | psc_score = torch.zeros(labels.shape) 137 | scale_count = 0 138 | for layer_index in range(self.start_index, self.start_index + self.n): 139 | layers = self.sorted_indices[:layer_index+1] 140 | # print(f'layers: {layers}') 141 | smodel = self.scale_var_index(layers, scale=self.scale) 142 | scale_count += 1 143 | smodel.eval() 144 | logits = smodel(imgs).detach().cpu() 145 | softmax_logits = torch.nn.functional.softmax(logits, dim=1) 146 | psc_score += softmax_logits[torch.arange(softmax_logits.size(0)), original_pred] 147 | 148 | psc_score /= scale_count 149 | all_psc_score.append(psc_score) 150 | 151 | all_psc_score = torch.cat(all_psc_score, dim=0) 152 | pred_correct_mask = torch.cat(pred_correct_mask, dim=0) 153 | all_psc_score = all_psc_score[pred_correct_mask] 154 | return all_psc_score 155 | def test(self, testset, poisoned_testset): 156 | print(f'start_index: {self.start_index}') 157 | 158 | benign_psc = self._test(testset) 159 | poison_psc = self._test(poisoned_testset) 160 | 161 | num_benign = benign_psc.size(0) 162 | num_poison = poison_psc.size(0) 163 | 164 | y_true = torch.cat((torch.zeros_like(benign_psc), torch.ones_like(poison_psc))) 165 | y_score = torch.cat((benign_psc, poison_psc), dim=0) 166 | y_pred = (y_score >= self.T) 167 | fpr, tpr, thresholds = metrics.roc_curve(y_true, y_score) 168 | auc = metrics.auc(fpr, tpr) 169 | tn, fp, fn, tp = metrics.confusion_matrix(y_true, y_pred).ravel() 170 | myf1 = metrics.f1_score(y_true, y_pred) 171 | print("TPR: {:.2f}".format(tp / (tp + fn) * 100)) 172 | print("FPR: {:.2f}".format(fp / (tn + fp) * 100)) 173 | print("AUC: {:.4f}".format(auc)) 174 | print(f"f1 score: {myf1}") 175 | 176 | def _detect(self, inputs): 177 | inputs = inputs.cuda() 178 | self.model.eval() 179 | self.model.cuda() 180 | original_pred = torch.argmax(self.model(inputs), dim=1) # model prediction 181 | 182 | psc_score = torch.zeros(inputs.size(0)) 183 | scale_count = 0 184 | for layer_index in range(self.start_index, self.start_index + self.n): 185 | layers = self.sorted_indices[:layer_index+1] 186 | # print(f'layers: {layers}') 187 | smodel = self.scale_var_index(layers, scale=self.scale) 188 | scale_count += 1 189 | smodel.eval() 190 | logits = smodel(inputs).detach().cpu() 191 | softmax_logits = torch.nn.functional.softmax(logits, dim=1) 192 | psc_score += softmax_logits[torch.arange(softmax_logits.size(0)), original_pred] 193 | 194 | psc_score /= scale_count 195 | 196 | y_pred = psc_score >= self.T 197 | return y_pred 198 | 199 | def detect(self, dataset): 200 | data_loader = torch.utils.data.DataLoader(dataset, batch_size=128, shuffle=False) 201 | with torch.no_grad(): 202 | for idx, batch in enumerate(data_loader): 203 | imgs = batch[0] 204 | return self._detect(imgs) 205 | -------------------------------------------------------------------------------- /core/defenses/Pruning.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the implement of pruning proposed in [1]. 3 | [1] Fine-Pruning: Defending Against Backdooring Attacks on Deep Neural Networks. RAID, 2018. 4 | ''' 5 | 6 | import os 7 | import torch 8 | import torch.nn as nn 9 | 10 | 11 | from .base import Base 12 | from ..utils import test 13 | from torch.utils.data import DataLoader 14 | 15 | 16 | # Define model pruning 17 | class MaskedLayer(nn.Module): 18 | def __init__(self, base, mask): 19 | super(MaskedLayer, self).__init__() 20 | self.base = base 21 | self.mask = mask 22 | 23 | def forward(self, input): 24 | return self.base(input) * self.mask 25 | 26 | 27 | 28 | class Pruning(Base): 29 | """Pruning process. 30 | Args: 31 | train_dataset (types in support_list): forward dataset. 32 | test_dataset (types in support_list): testing dataset. 33 | model (torch.nn.Module): Network. 34 | layer(list): The layers to prune 35 | prune_rate (double): the pruning rate 36 | schedule (dict): Training or testing schedule. Default: None. 37 | seed (int): Global seed for random numbers. Default: 0. 38 | deterministic (bool): Sets whether PyTorch operations must use "deterministic" algorithms. 39 | That is, algorithms which, given the same input, and when run on the same software and hardware, 40 | always produce the same output. When enabled, operations will use deterministic algorithms when available, 41 | and if only nondeterministic algorithms are available they will throw a RuntimeError when called. Default: False. 42 | """ 43 | 44 | def __init__(self, 45 | train_dataset=None, 46 | test_dataset=None, 47 | model=None, 48 | layer=None, 49 | prune_rate=None, 50 | schedule=None, 51 | seed=0, 52 | deterministic=False): 53 | super(Pruning, self).__init__(seed=seed, deterministic=deterministic) 54 | 55 | self.train_dataset = train_dataset 56 | self.test_dataset = test_dataset 57 | self.model = model 58 | self.layer = layer 59 | self.prune_rate = prune_rate 60 | self.schedule = schedule 61 | 62 | 63 | def repair(self, schedule=None): 64 | """pruning. 65 | Args: 66 | schedule (dict): Schedule for testing. 67 | """ 68 | 69 | if schedule == None: 70 | raise AttributeError("Schedule is None, please check your schedule setting.") 71 | current_schedule = schedule 72 | 73 | 74 | # Use GPU 75 | if 'device' in current_schedule and current_schedule['device'] == 'GPU': 76 | if 'CUDA_VISIBLE_DEVICES' in current_schedule: 77 | os.environ['CUDA_VISIBLE_DEVICES'] = current_schedule['CUDA_VISIBLE_DEVICES'] 78 | 79 | assert torch.cuda.device_count() > 0, 'This machine has no cuda devices!' 80 | assert current_schedule['GPU_num'] > 0, 'GPU_num should be a positive integer' 81 | print( 82 | f"This machine has {torch.cuda.device_count()} cuda devices, and use {current_schedule['GPU_num']} of them to train.") 83 | 84 | if current_schedule['GPU_num'] == 1: 85 | device = torch.device("cuda:0") 86 | else: 87 | gpus = list(range(current_schedule['GPU_num'])) 88 | self.model = nn.DataParallel(self.model.cuda(), device_ids=gpus, output_device=gpus[0]) 89 | # TODO: DDP training 90 | pass 91 | # Use CPU 92 | else: 93 | device = torch.device("cpu") 94 | 95 | model = self.model.to(device) 96 | layer_to_prune = self.layer 97 | tr_loader = DataLoader(self.train_dataset, batch_size=current_schedule['batch_size'], 98 | num_workers=current_schedule['num_workers'], 99 | drop_last=True, pin_memory=True) 100 | prune_rate = self.prune_rate 101 | 102 | 103 | # prune silent activation 104 | print("======== pruning... ========") 105 | with torch.no_grad(): 106 | container = [] 107 | 108 | def forward_hook(module, input, output): 109 | container.append(output) 110 | 111 | hook = getattr(model, layer_to_prune).register_forward_hook(forward_hook) 112 | print("Forwarding all training set") 113 | 114 | model.eval() 115 | for data, _ in tr_loader: 116 | model(data.cuda()) 117 | hook.remove() 118 | 119 | container = torch.cat(container, dim=0) 120 | activation = torch.mean(container, dim=[0, 2, 3]) 121 | seq_sort = torch.argsort(activation) 122 | num_channels = len(activation) 123 | prunned_channels = int(num_channels * prune_rate) 124 | mask = torch.ones(num_channels).cuda() 125 | for element in seq_sort[:prunned_channels]: 126 | mask[element] = 0 127 | if len(container.shape) == 4: 128 | mask = mask.reshape(1, -1, 1, 1) 129 | setattr(model, layer_to_prune, MaskedLayer(getattr(model, layer_to_prune), mask)) 130 | 131 | self.model = model 132 | print("======== pruning complete ========") 133 | 134 | 135 | def test(self, schedule=None): 136 | """Test the pruned model. 137 | Args: 138 | schedule (dict): Schedule for testing. 139 | """ 140 | if schedule == None: 141 | raise AttributeError("Schedule is None, please check your schedule setting.") 142 | if self.test_dataset == None: 143 | raise AttributeError("Test set is None, please check your setting.") 144 | test(self.model, self.test_dataset, schedule) 145 | 146 | def get_model(self): 147 | return self.model 148 | -------------------------------------------------------------------------------- /core/defenses/SCALE_UP.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import pdb 4 | import torch 5 | from torchvision import transforms 6 | from sklearn import metrics 7 | from tqdm import tqdm 8 | import copy 9 | import numpy as np 10 | import torch.nn.functional as F 11 | import umap 12 | from PIL import Image 13 | import torchvision 14 | import matplotlib.pyplot as plt 15 | from collections import Counter 16 | from torch.utils.data import Subset 17 | import torch.nn as nn 18 | import matplotlib.pyplot as plt 19 | import numpy as np 20 | 21 | from ..utils import test 22 | 23 | from .base import Base 24 | 25 | 26 | 27 | class SCALE_UP(Base): 28 | """Identify and filter malicious testing samples (SCALE-UP). 29 | 30 | Args: 31 | model (nn.Module): The original backdoored model. 32 | scale_set (List): The hyper-parameter for a set of scaling factors. Each integer n in the set scales the pixel values of an input image "x" by a factor of n. 33 | T (float): The hyper-parameter for defender-specified threshold T. If SPC(x) > T , we deem it as a backdoor sample. 34 | valset (Dataset): In data-limited scaled prediction consistency analysis, we assume that defenders have a few benign samples from each class. 35 | seed (int): Global seed for random numbers. Default: 0. 36 | deterministic (bool): Sets whether PyTorch operations must use "deterministic" algorithms. 37 | 38 | """ 39 | def __init__(self, model, scale_set=None, threshold=None, valset=None, seed=666, deterministic=False): 40 | super(SCALE_UP, self).__init__(seed, deterministic) 41 | self.model = model 42 | self.model.cuda() 43 | self.model.eval() 44 | if scale_set is None: 45 | scale_set = [3, 5, 7, 9, 11] 46 | if threshold is None: 47 | self.T = 0.5 48 | self.scale_set = scale_set 49 | # The statics for SPC values on samples from different classes 50 | self.valset = valset 51 | if self.valset: 52 | self.mean = None 53 | self.std = None 54 | self.init_spc_norm(self.valset) 55 | 56 | # test accuracy on the dataset 57 | def test_acc(self, dataset, schedule): 58 | """Test repaired curve model on dataset 59 | 60 | Args: 61 | dataset (types in support_list): Dataset. 62 | schedule (dict): Schedule for testing. 63 | """ 64 | model = self.model 65 | test(model, dataset, schedule) 66 | 67 | def init_spc_norm(self, valset): 68 | val_loader = torch.utils.data.DataLoader(valset, batch_size=128, shuffle=False) 69 | total_spc = [] 70 | for idx, batch in enumerate(val_loader): 71 | clean_img = batch[0] 72 | labels = batch[1] 73 | clean_img = clean_img.cuda() # batch * channels * hight * width 74 | labels = labels.cuda() # batch 75 | scaled_imgs = [] 76 | scaled_labels = [] 77 | for scale in self.scale_set: 78 | # If normalize: 79 | # scaled_imgs.append(self.normalizer(torch.clip(self.denormalizer(clean_img) * scale, 0.0, 1.0))) 80 | # No normalize 81 | scaled_imgs.append(torch.clip(clean_img * scale, 0.0, 1.0)) 82 | for scale_img in scaled_imgs: 83 | scale_label = torch.argmax(self.model(scale_img), dim=1) 84 | scaled_labels.append(scale_label) 85 | 86 | # compute the SPC Value 87 | spc = torch.zeros(labels.shape).cuda() 88 | for scale_label in scaled_labels: 89 | spc += scale_label == labels 90 | spc /= len(self.scale_set) 91 | total_spc.append(spc) 92 | total_spc = torch.cat(total_spc) 93 | 94 | self.mean = torch.mean(total_spc).item() 95 | self.std = torch.std(total_spc).item() 96 | 97 | def _test(self, dataset): 98 | data_loader = torch.utils.data.DataLoader(dataset, batch_size=128, shuffle=False) 99 | self.model.eval() 100 | all_spc_score = [] 101 | pred_correct_mask = [] 102 | 103 | with torch.no_grad(): 104 | for idx, batch in enumerate(data_loader): 105 | imgs = batch[0] 106 | labels = batch[1] 107 | imgs = imgs.cuda() # batch * channels * hight * width 108 | labels = labels.cuda() # batch 109 | original_pred = torch.argmax(self.model(imgs), dim=1) # model prediction 110 | mask = torch.eq(labels, original_pred) # only look at those samples that successfully attack the DNN 111 | pred_correct_mask.append(mask) 112 | 113 | scaled_imgs = [] 114 | scaled_labels = [] 115 | for scale in self.scale_set: 116 | scaled_imgs.append(torch.clip(imgs * scale, 0.0, 1.0)) 117 | # normalized 118 | # scaled_imgs.append(self.normalizer(torch.clip(self.denormalizer(clean_img) * scale, 0.0, 1.0))) 119 | for scale_img in scaled_imgs: 120 | scale_label = torch.argmax(self.model(scale_img), dim=1) 121 | scaled_labels.append(scale_label) 122 | 123 | spc_score = torch.zeros(labels.shape).cuda() 124 | for scale_label in scaled_labels: 125 | spc_score += scale_label == original_pred 126 | spc_score /= len(self.scale_set) 127 | 128 | if self.valset: 129 | spc_score = (spc_score - self.mean) / self.std 130 | all_spc_score.append(spc_score) 131 | 132 | all_spc_score = torch.cat(all_spc_score, dim=0).cpu() 133 | pred_correct_mask = torch.cat(pred_correct_mask, dim=0) 134 | all_spc_score = all_spc_score[pred_correct_mask] 135 | return all_spc_score 136 | 137 | def test(self, testset, poisoned_testset): 138 | 139 | benign_spc = self._test(testset) 140 | poison_spc = self._test(poisoned_testset) 141 | 142 | num_benign = benign_spc.size(0) 143 | num_poison = poison_spc.size(0) 144 | 145 | y_true = torch.cat((torch.zeros_like(benign_spc), torch.ones_like(poison_spc))) 146 | y_score = torch.cat((benign_spc, poison_spc), dim=0) 147 | y_pred = (y_score >= self.T) 148 | fpr, tpr, thresholds = metrics.roc_curve(y_true, y_score) 149 | auc = metrics.auc(fpr, tpr) 150 | tn, fp, fn, tp = metrics.confusion_matrix(y_true, y_pred).ravel() 151 | myf1 = metrics.f1_score(y_true, y_pred) 152 | print("TPR: {:.2f}".format(tp / (tp + fn) * 100)) 153 | print("FPR: {:.2f}".format(fp / (tn + fp) * 100)) 154 | print("AUC: {:.4f}".format(auc)) 155 | print(f"f1 score: {myf1}") 156 | 157 | # inputs: Tensor List , e.g., tensor([tensor_img1, tensor_img2, ..., tensor_imgn]) 158 | def _detect(self, inputs): 159 | inputs = inputs.cuda() 160 | self.model.eval() 161 | self.model.cuda() 162 | original_pred = torch.argmax(self.model(inputs), dim=1) # model prediction 163 | 164 | scaled_imgs = [] 165 | scaled_labels = [] 166 | for scale in self.scale_set: 167 | scaled_imgs.append(torch.clip(inputs * scale, 0.0, 1.0)) 168 | # normalized 169 | # scaled_imgs.append(self.normalizer(torch.clip(self.denormalizer(clean_img) * scale, 0.0, 1.0))) 170 | for scale_img in scaled_imgs: 171 | scale_label = torch.argmax(self.model(scale_img), dim=1) 172 | scaled_labels.append(scale_label) 173 | 174 | spc_score = torch.zeros(inputs.size(0)).cuda() 175 | for scale_label in scaled_labels: 176 | spc_score += scale_label == original_pred 177 | spc_score /= len(self.scale_set) 178 | 179 | if self.valset: 180 | spc_score = (spc_score - self.mean) / self.std 181 | 182 | y_pred = (spc_score >= self.T) 183 | return y_pred 184 | 185 | 186 | def detect(self, dataset): 187 | data_loader = torch.utils.data.DataLoader(dataset, batch_size=128, shuffle=False) 188 | with torch.no_grad(): 189 | for idx, batch in enumerate(data_loader): 190 | imgs = batch[0] 191 | return self._detect(imgs) 192 | 193 | -------------------------------------------------------------------------------- /core/defenses/ShrinkPad.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the implement of pre-processing-based backdoor defense with ShrinkPad proposed in [1]. 3 | 4 | Reference: 5 | [1] Backdoor Attack in the Physical World. ICLR Workshop, 2021. 6 | ''' 7 | 8 | 9 | import os 10 | from copy import deepcopy 11 | 12 | import torch 13 | import torchvision.transforms as transforms 14 | 15 | from .base import Base 16 | from ..utils import test 17 | 18 | 19 | def RandomPad(sum_w, sum_h, fill=0): 20 | transforms_bag=[] 21 | for i in range(sum_w+1): 22 | for j in range(sum_h+1): 23 | transforms_bag.append(transforms.Pad(padding=(i,j,sum_w-i,sum_h-j))) 24 | 25 | return transforms_bag 26 | 27 | 28 | def build_ShrinkPad(size_map, pad): 29 | return transforms.Compose([ 30 | transforms.Resize((size_map - pad, size_map - pad)), 31 | transforms.RandomChoice(RandomPad(sum_w=pad, sum_h=pad)) 32 | ]) 33 | 34 | 35 | class ShrinkPad(Base): 36 | """Construct defense datasets with ShrinkPad method. 37 | 38 | Args: 39 | size_map (int): Size of image. 40 | pad (int): Size of pad. 41 | seed (int): Global seed for random numbers. Default: 0. 42 | deterministic (bool): Sets whether PyTorch operations must use "deterministic" algorithms. 43 | That is, algorithms which, given the same input, and when run on the same software and hardware, 44 | always produce the same output. When enabled, operations will use deterministic algorithms when available, 45 | and if only nondeterministic algorithms are available they will throw a RuntimeError when called. Default: False. 46 | """ 47 | 48 | def __init__(self, 49 | size_map, 50 | pad, 51 | seed=0, 52 | deterministic=False): 53 | 54 | super(ShrinkPad, self).__init__(seed=seed, deterministic=deterministic) 55 | 56 | self.global_size_map = size_map 57 | self.current_size_map = None 58 | 59 | self.global_pad = pad 60 | self.current_pad = None 61 | 62 | def preprocess(self, data, size_map=None, pad=None): 63 | """Perform ShrinkPad defense method on data and return the preprocessed data. 64 | 65 | Args: 66 | data (torch.Tensor): Input data. 67 | size_map (int): Size of image. Default: None. 68 | pad (int): Size of pad. Default: None. 69 | 70 | Returns: 71 | torch.Tensor: The preprocessed data. 72 | """ 73 | if size_map is None: 74 | self.current_size_map = self.global_size_map 75 | else: 76 | self.current_size_map = size_map 77 | 78 | if pad is None: 79 | self.current_pad = self.global_pad 80 | else: 81 | self.current_pad = pad 82 | 83 | shrinkpad = build_ShrinkPad(self.current_size_map, self.current_pad) 84 | return shrinkpad(data) 85 | 86 | def _predict(self, model, data, device, batch_size, num_workers): 87 | with torch.no_grad(): 88 | model = model.to(device) 89 | model.eval() 90 | 91 | predict_digits = [] 92 | for i in range(data.shape[0] // batch_size): 93 | # breakpoint() 94 | batch_img = data[i*batch_size:(i+1)*batch_size, ...] 95 | batch_img = batch_img.to(device) 96 | batch_img = model(batch_img) 97 | batch_img = batch_img.cpu() 98 | predict_digits.append(batch_img) 99 | 100 | if data.shape[0] % batch_size != 0: 101 | batch_img = data[(data.shape[0] // batch_size) * batch_size:, ...] 102 | batch_img = batch_img.to(device) 103 | batch_img = model(batch_img) 104 | batch_img = batch_img.cpu() 105 | predict_digits.append(batch_img) 106 | 107 | predict_digits = torch.cat(predict_digits, dim=0) 108 | return predict_digits 109 | 110 | def predict(self, model, data, schedule, size_map=None, pad=None): 111 | """Apply ShrinkPad defense method to input data and get the predicts. 112 | 113 | Args: 114 | model (torch.nn.Module): Network. 115 | data (torch.Tensor): Input data. 116 | schedule (dict): Schedule for predicting. 117 | size_map (int): Size of image. Default: None. 118 | pad (int): Size of pad. Default: None. 119 | 120 | Returns: 121 | torch.Tensor: The predicts. 122 | """ 123 | if size_map is None: 124 | self.current_size_map = self.global_size_map 125 | else: 126 | self.current_size_map = size_map 127 | 128 | if pad is None: 129 | self.current_pad = self.global_pad 130 | else: 131 | self.current_pad = pad 132 | 133 | shrinkpad = build_ShrinkPad(self.current_size_map, self.current_pad) 134 | preprocessed_data = self.preprocess(data) 135 | 136 | if 'test_model' in schedule: 137 | model.load_state_dict(torch.load(schedule['test_model']), strict=False) 138 | 139 | # Use GPU 140 | if 'device' in schedule and schedule['device'] == 'GPU': 141 | if 'CUDA_VISIBLE_DEVICES' in schedule: 142 | os.environ['CUDA_VISIBLE_DEVICES'] = schedule['CUDA_VISIBLE_DEVICES'] 143 | 144 | assert torch.cuda.device_count() > 0, 'This machine has no cuda devices!' 145 | assert schedule['GPU_num'] >0, 'GPU_num should be a positive integer' 146 | print(f"This machine has {torch.cuda.device_count()} cuda devices, and use {schedule['GPU_num']} of them to train.") 147 | 148 | if schedule['GPU_num'] == 1: 149 | device = torch.device("cuda:0") 150 | else: 151 | gpus = list(range(schedule['GPU_num'])) 152 | model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0]) 153 | # TODO: DDP training 154 | pass 155 | # Use CPU 156 | else: 157 | device = torch.device("cpu") 158 | 159 | return self._predict(model, preprocessed_data, device, schedule['batch_size'], schedule['num_workers']) 160 | 161 | def test(self, model, dataset, schedule, size_map=None, pad=None): 162 | """Test ShrinkPad on dataset. 163 | 164 | Args: 165 | model (torch.nn.Module): Network. 166 | dataset (types in support_list): Dataset. 167 | schedule (dict): Schedule for testing. 168 | size_map (int): Size of image. Default: None. 169 | pad (int): Size of pad. Default: None. 170 | """ 171 | if size_map is None: 172 | self.current_size_map = self.global_size_map 173 | else: 174 | self.current_size_map = size_map 175 | 176 | if pad is None: 177 | self.current_pad = self.global_pad 178 | else: 179 | self.current_pad = pad 180 | 181 | defense_dataset = deepcopy(dataset) 182 | # defense_dataset.transform.transforms.append(transforms.ToPILImage()) 183 | defense_dataset.transform.transforms.append(build_ShrinkPad(self.current_size_map, self.current_pad)) 184 | # defense_dataset.transform.transforms.append(transforms.ToTensor()) 185 | 186 | if hasattr(defense_dataset, 'poisoned_transform'): 187 | # defense_dataset.poisoned_transform.transforms.append(transforms.ToPILImage()) 188 | defense_dataset.poisoned_transform.transforms.append(build_ShrinkPad(self.current_size_map, self.current_pad)) 189 | # defense_dataset.poisoned_transform.transforms.append(transforms.ToTensor()) 190 | 191 | test(model, defense_dataset, schedule) 192 | -------------------------------------------------------------------------------- /core/defenses/__init__.py: -------------------------------------------------------------------------------- 1 | from .ABL import ABL 2 | from .AutoEncoderDefense import AutoEncoderDefense 3 | from .ShrinkPad import ShrinkPad 4 | from .MCR import MCR 5 | from .FineTuning import FineTuning 6 | from .NAD import NAD 7 | from .Pruning import Pruning 8 | from .CutMix import CutMix 9 | from .IBD_PSC import IBD_PSC 10 | from .SCALE_UP import SCALE_UP 11 | from .REFINE import REFINE 12 | 13 | __all__ = [ 14 | 'AutoEncoderDefense', 'ShrinkPad', 'FineTuning', 'MCR', 'NAD', 'Pruning', 'ABL', 'CutMix', 'IBD_PSC', 'SCALE_UP', 'REFINE' 15 | ] 16 | -------------------------------------------------------------------------------- /core/defenses/base.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | 4 | import torch 5 | import numpy as np 6 | 7 | 8 | class Base(object): 9 | """Base class for backdoor defense. 10 | 11 | Args: 12 | seed (int): Global seed for random numbers. Default: 0. 13 | deterministic (bool): Sets whether PyTorch operations must use "deterministic" algorithms. 14 | That is, algorithms which, given the same input, and when run on the same software and hardware, 15 | always produce the same output. When enabled, operations will use deterministic algorithms when available, 16 | and if only nondeterministic algorithms are available they will throw a RuntimeError when called. Default: False. 17 | """ 18 | 19 | def __init__(self, seed=0, deterministic=False): 20 | self._set_seed(seed, deterministic) 21 | 22 | def _set_seed(self, seed, deterministic): 23 | # Use torch.manual_seed() to seed the RNG for all devices (both CPU and CUDA). 24 | torch.manual_seed(seed) 25 | 26 | # Set python seed 27 | random.seed(seed) 28 | 29 | # Set numpy seed (However, some applications and libraries may use NumPy Random Generator objects, 30 | # not the global RNG (https://numpy.org/doc/stable/reference/random/generator.html), and those will 31 | # need to be seeded consistently as well.) 32 | np.random.seed(seed) 33 | 34 | os.environ['PYTHONHASHSEED'] = str(seed) 35 | 36 | if deterministic: 37 | torch.backends.cudnn.benchmark = False 38 | torch.use_deterministic_algorithms(True) 39 | torch.backends.cudnn.deterministic = True 40 | os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' 41 | # Hint: In some versions of CUDA, RNNs and LSTM networks may have non-deterministic behavior. 42 | # If you want to set them deterministic, see torch.nn.RNN() and torch.nn.LSTM() for details and workarounds. 43 | -------------------------------------------------------------------------------- /core/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .autoencoder import AutoEncoder 2 | from .baseline_MNIST_network import BaselineMNISTNetwork 3 | from .resnet import ResNet 4 | from .vgg import * 5 | from .unet import UNet, UNetLittle 6 | 7 | __all__ = [ 8 | 'AutoEncoder', 'BaselineMNISTNetwork', 'ResNet', 'UNet', 'UNetLittle' 9 | ] -------------------------------------------------------------------------------- /core/models/autoencoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class AutoEncoder1x28x28(nn.Module): 6 | """Autoencoder for 1x28x28 input image. 7 | 8 | This is a reimplementation of the blog post 'Building Autoencoders in Keras', from blog 9 | `Building Autoencoders in Keras `_. 10 | """ 11 | def __init__(self): 12 | super(AutoEncoder1x28x28, self).__init__() 13 | self.conv1 = nn.Conv2d(1, 32, 3, padding=1) 14 | self.conv2 = nn.Conv2d(32, 32, 3, padding=1) 15 | self.conv3 = nn.Conv2d(32, 32, 3, padding=1) 16 | self.conv4 = nn.Conv2d(32, 32, 3, padding=1) 17 | self.conv5 = nn.Conv2d(32, 1, 3, padding=1) 18 | 19 | self.relu = nn.ReLU(inplace=True) 20 | self.max_pool = nn.MaxPool2d(2) 21 | self.upsampler = nn.UpsamplingNearest2d(scale_factor=2.0) 22 | self.sigmoid = nn.Sigmoid() 23 | 24 | def forward(self, x): 25 | # Input size: [batch, 1, 28, 28] 26 | # Output size: [batch, 1, 28, 28] 27 | x = self.conv1(x) # [batch, 32, 28, 28] 28 | x = self.relu(x) # [batch, 32, 28, 28] 29 | x = self.max_pool(x) # [batch, 32, 14, 14] 30 | 31 | x = self.conv2(x) # [batch, 32, 14, 14] 32 | x = self.relu(x) # [batch, 32, 14, 14] 33 | x = self.max_pool(x) # [batch, 32, 7, 7] 34 | 35 | x = self.conv3(x) # [batch, 32, 7, 7] 36 | x = self.relu(x) # [batch, 32, 7, 7] 37 | x = self.upsampler(x) # [batch, 32, 14, 14] 38 | 39 | x = self.conv4(x) # [batch, 32, 14, 14] 40 | x = self.relu(x) # [batch, 32, 14, 14] 41 | x = self.upsampler(x) # [batch, 32, 28, 28] 42 | 43 | x = self.conv5(x) # [batch, 1, 28, 28] 44 | x = self.sigmoid(x) # [batch, 1, 28, 28] 45 | 46 | return x 47 | 48 | 49 | class AutoEncoder3x32x32(nn.Module): 50 | """Autoencoder for 3x32x32 input image. 51 | 52 | This is modified from 'PyTorch-CIFAR-10-autoencoder', from github 53 | `PyTorch-CIFAR-10-autoencoder `_. 54 | """ 55 | def __init__(self): 56 | super(AutoEncoder3x32x32, self).__init__() 57 | # Input size: [batch, 3, 32, 32] 58 | # Output size: [batch, 3, 32, 32] 59 | self.encoder = nn.Sequential( 60 | nn.Conv2d(3, 12, 4, stride=2, padding=1), # [batch, 12, 16, 16] 61 | nn.ReLU(), 62 | nn.Conv2d(12, 24, 4, stride=2, padding=1), # [batch, 24, 8, 8] 63 | nn.ReLU(), 64 | nn.Conv2d(24, 48, 4, stride=2, padding=1), # [batch, 48, 4, 4] 65 | nn.ReLU(), 66 | # nn.Conv2d(48, 96, 4, stride=2, padding=1), # [batch, 96, 2, 2] 67 | # nn.ReLU(), 68 | ) 69 | self.decoder = nn.Sequential( 70 | # nn.ConvTranspose2d(96, 48, 4, stride=2, padding=1), # [batch, 48, 4, 4] 71 | # nn.ReLU(), 72 | nn.ConvTranspose2d(48, 24, 4, stride=2, padding=1), # [batch, 24, 8, 8] 73 | nn.ReLU(), 74 | nn.ConvTranspose2d(24, 12, 4, stride=2, padding=1), # [batch, 12, 16, 16] 75 | nn.ReLU(), 76 | nn.ConvTranspose2d(12, 3, 4, stride=2, padding=1), # [batch, 3, 32, 32] 77 | nn.Sigmoid(), 78 | ) 79 | 80 | def forward(self, x): 81 | encoded = self.encoder(x) 82 | decoded = self.decoder(encoded) 83 | return decoded 84 | 85 | 86 | def AutoEncoder(img_size): 87 | assert isinstance(img_size, tuple), f"img_size should be tuple, but got {type(img_size)}" 88 | if img_size == (1, 28, 28): 89 | return AutoEncoder1x28x28() 90 | elif img_size == (3, 32, 32): 91 | return AutoEncoder3x32x32() 92 | else: 93 | raise NotImplementedError("Unsupported img size!") 94 | 95 | 96 | if __name__=='__main__': 97 | x = torch.randn((128, 1, 28, 28)) 98 | # model = AutoEncoder([1, 28, 28]) 99 | model = AutoEncoder((1, 28, 28)) 100 | x = model(x) 101 | 102 | x = torch.randn((128, 3, 32, 32)) 103 | model = AutoEncoder((3, 32, 32)) 104 | x = model(x) 105 | -------------------------------------------------------------------------------- /core/models/baseline_MNIST_network.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class BaselineMNISTNetwork(nn.Module): 6 | """Baseline network for MNIST dataset. 7 | 8 | This network is the implement of baseline network for MNIST dataset, from paper 9 | `BadNets: Evaluating Backdooring Attackson Deep Neural Networks `_. 10 | """ 11 | 12 | def __init__(self): 13 | super(BaselineMNISTNetwork, self).__init__() 14 | self.conv1 = nn.Conv2d(1, 16, 5) 15 | self.conv2 = nn.Conv2d(16, 32, 5) 16 | self.fc1 = nn.Linear(512, 512) 17 | self.fc2 = nn.Linear(512, 10) 18 | 19 | self.avg_pool = nn.AvgPool2d(2) 20 | self.relu = nn.ReLU(inplace=True) 21 | # self.softmax = nn.Softmax(dim=1) 22 | 23 | def forward(self, x): 24 | x = self.conv1(x) 25 | x = self.relu(x) 26 | x = self.avg_pool(x) 27 | 28 | x = self.conv2(x) 29 | x = self.relu(x) 30 | x = self.avg_pool(x) 31 | 32 | x = x.contiguous().view(-1, 512) 33 | 34 | x = self.fc1(x) 35 | x = self.relu(x) 36 | 37 | x = self.fc2(x) 38 | # x = self.softmax(x) 39 | 40 | return x 41 | 42 | if __name__ == '__main__': 43 | baseline_MNIST_network = BaselineMNISTNetwork() 44 | x = torch.randn(16, 1, 28, 28) 45 | x = baseline_MNIST_network(x) 46 | print(x.size()) 47 | print(x) 48 | -------------------------------------------------------------------------------- /core/models/resnet.py: -------------------------------------------------------------------------------- 1 | ''' 2 | ResNet in PyTorch. 3 | For Pre-activation ResNet, see 'preact_resnet.py'. 4 | Reference: 5 | [1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun 6 | Deep Residual Learning for Image Recognition. arXiv:1512.03385 7 | ''' 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | 11 | 12 | class BasicBlock(nn.Module): 13 | expansion = 1 14 | 15 | def __init__(self, in_planes, planes, stride=1): 16 | super(BasicBlock, self).__init__() 17 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 18 | self.bn1 = nn.BatchNorm2d(planes) 19 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) 20 | self.bn2 = nn.BatchNorm2d(planes) 21 | 22 | self.shortcut = nn.Sequential() 23 | if stride != 1 or in_planes != self.expansion*planes: 24 | self.shortcut = nn.Sequential( 25 | nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False), 26 | nn.BatchNorm2d(self.expansion*planes) 27 | ) 28 | 29 | def forward(self, x): 30 | out = F.relu(self.bn1(self.conv1(x))) 31 | out = self.bn2(self.conv2(out)) 32 | out += self.shortcut(x) 33 | out = F.relu(out) 34 | return out 35 | 36 | 37 | class Bottleneck(nn.Module): 38 | expansion = 4 39 | 40 | def __init__(self, in_planes, planes, stride=1): 41 | super(Bottleneck, self).__init__() 42 | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False) 43 | self.bn1 = nn.BatchNorm2d(planes) 44 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 45 | self.bn2 = nn.BatchNorm2d(planes) 46 | self.conv3 = nn.Conv2d(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, num_classes=10): 67 | super(_ResNet, self).__init__() 68 | self.in_planes = 64 69 | 70 | self.conv1 = nn.Conv2d(3, 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.linear = nn.Linear(512*block.expansion, num_classes) 77 | 78 | def _make_layer(self, block, planes, num_blocks, stride): 79 | strides = [stride] + [1]*(num_blocks-1) 80 | layers = [] 81 | for stride in strides: 82 | layers.append(block(self.in_planes, planes, stride)) 83 | self.in_planes = planes * block.expansion 84 | return nn.Sequential(*layers) 85 | 86 | def forward(self, x): 87 | out = F.relu(self.bn1(self.conv1(x))) 88 | out = self.layer1(out) 89 | out = self.layer2(out) 90 | out = self.layer3(out) 91 | out = self.layer4(out) 92 | out = F.avg_pool2d(out, 4) 93 | out = out.view(out.size(0), -1) 94 | out = self.linear(out) 95 | return out 96 | 97 | 98 | def ResNet(num, num_classes=10): 99 | if num == 18: 100 | return _ResNet(BasicBlock, [2,2,2,2], num_classes) 101 | elif num == 34: 102 | return _ResNet(BasicBlock, [3,4,6,3], num_classes) 103 | elif num == 50: 104 | return _ResNet(Bottleneck, [3,4,6,3], num_classes) 105 | elif num == 101: 106 | return _ResNet(Bottleneck, [3,4,23,3], num_classes) 107 | elif num == 152: 108 | return _ResNet(Bottleneck, [3,8,36,3], num_classes) 109 | else: 110 | raise NotImplementedError 111 | -------------------------------------------------------------------------------- /core/models/resnet_curve.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the part of the implement of model-repairing-based backdoor defense with MCR proposed in [1]. 3 | 4 | Reference: 5 | [1] Bridging Mode Connectivity in Loss Landscapes and Adversarial Robustness. ICLR, 2020. 6 | ''' 7 | from calendar import c 8 | import sys, os 9 | sys.path.append(os.path.join(os.getcwd(), 'core/models')) 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | import curves 13 | import math 14 | 15 | class BasicBlockCurves(nn.Module): 16 | expansion = 1 17 | 18 | def __init__(self, in_planes, planes, fix_points, stride=1): 19 | super(BasicBlockCurves, self).__init__() 20 | self.conv1 = curves.Conv2d(in_planes, planes, kernel_size=3, fix_points=fix_points, stride=stride, padding=1, bias=False) 21 | self.bn1 = curves.BatchNorm2d(planes, fix_points) 22 | self.conv2 = curves.Conv2d(planes, planes, kernel_size=3, fix_points=fix_points, stride=1, padding=1, bias=False) 23 | self.bn2 = curves.BatchNorm2d(planes, fix_points) 24 | 25 | self.shortcut = nn.Sequential() 26 | if stride != 1 or in_planes != self.expansion*planes: 27 | self.shortcut = nn.Sequential( 28 | curves.Conv2d(in_planes, self.expansion*planes, kernel_size=1, fix_points=fix_points, stride=stride, bias=False), 29 | curves.BatchNorm2d(self.expansion*planes, fix_points) 30 | ) 31 | 32 | def forward(self, x, coeffs_t): 33 | out = F.relu(self.bn1(self.conv1(x, coeffs_t), coeffs_t)) 34 | out = self.bn2(self.conv2(out, coeffs_t), coeffs_t) 35 | # out += self.shortcut(x) 36 | for module in self.shortcut: 37 | if isinstance(module, curves.CurveModule): 38 | x = module(x, coeffs_t) 39 | else: 40 | x = module(x) 41 | out += x 42 | out = F.relu(out) 43 | return out 44 | 45 | 46 | class BottleneckCurve(nn.Module): 47 | expansion = 4 48 | 49 | def __init__(self, in_planes, planes, fix_points, stride=1): 50 | super(BottleneckCurve, self).__init__() 51 | self.conv1 = curves.Conv2d(in_planes, planes, kernel_size=1, fix_points=fix_points, bias=False) 52 | self.bn1 = curves.BatchNorm2d(planes, fix_points) 53 | self.conv2 = curves.Conv2d(planes, planes, kernel_size=3, fix_points=fix_points, stride=stride, padding=1, bias=False) 54 | self.bn2 = curves.BatchNorm2d(planes, fix_points) 55 | self.conv3 = curves.Conv2d(planes, self.expansion*planes, fix_points=fix_points, kernel_size=1, bias=False) 56 | self.bn3 = curves.BatchNorm2d(self.expansion*planes, fix_points) 57 | 58 | self.shortcut = nn.Sequential() 59 | if stride != 1 or in_planes != self.expansion*planes: 60 | self.shortcut = nn.Sequential( 61 | curves.Conv2d(in_planes, self.expansion*planes, kernel_size=1, fix_points=fix_points, stride=stride, bias=False), 62 | curves.BatchNorm2d(self.expansion*planes, fix_points) 63 | ) 64 | 65 | def forward(self, x, coeffs_t): 66 | out = F.relu(self.bn1(self.conv1(x, coeffs_t), coeffs_t)) 67 | out = F.relu(self.bn2(self.conv2(out, coeffs_t), coeffs_t)) 68 | out = self.bn3(self.conv3(out, coeffs_t), coeffs_t) 69 | # out += self.shortcut(x) 70 | for module in self.shortcut: 71 | if isinstance(module, curves.CurveModule): 72 | x = module(x, coeffs_t) 73 | else: 74 | x = module(x) 75 | out += x 76 | out = F.relu(out) 77 | return out 78 | 79 | 80 | class _ResNetCurve(nn.Module): 81 | def __init__(self, block, num_blocks, fix_points, num_classes=10, initialize=False): 82 | super(_ResNetCurve, self).__init__() 83 | self.in_planes = 64 84 | 85 | self.conv1 = curves.Conv2d(3, 64, kernel_size=3, fix_points=fix_points, stride=1, padding=1, bias=False) 86 | self.bn1 = curves.BatchNorm2d(64, fix_points) 87 | self.layer1 = self._make_layer(block, 64, num_blocks[0], fix_points, stride=1) 88 | self.layer2 = self._make_layer(block, 128, num_blocks[1], fix_points, stride=2) 89 | self.layer3 = self._make_layer(block, 256, num_blocks[2], fix_points, stride=2) 90 | self.layer4 = self._make_layer(block, 512, num_blocks[3], fix_points, stride=2) 91 | self.linear = curves.Linear(512*block.expansion, num_classes, fix_points) 92 | 93 | if initialize: 94 | for m in self.modules(): 95 | if isinstance(m, curves.Conv2d): 96 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 97 | for i in range(m.num_bends): 98 | getattr(m, 'weight_%d' % i).data.normal_(0, math.sqrt(2. / n)) 99 | elif isinstance(m, curves.BatchNorm2d): 100 | for i in range(m.num_bends): 101 | getattr(m, 'weight_%d' % i).data.fill_(1) 102 | getattr(m, 'bias_%d' % i).data.zero_() 103 | 104 | def _make_layer(self, block, planes, num_blocks, fix_points, stride): 105 | strides = [stride] + [1]*(num_blocks-1) 106 | layers = [] 107 | for stride in strides: 108 | layers.append(block(self.in_planes, planes, fix_points, stride)) 109 | self.in_planes = planes * block.expansion 110 | return nn.Sequential(*layers) 111 | 112 | def forward(self, x, coeffs_t): 113 | out = F.relu(self.bn1(self.conv1(x, coeffs_t), coeffs_t)) 114 | # out = self.layer1(out, coeffs_t) 115 | # out = self.layer2(out, coeffs_t) 116 | # out = self.layer3(out, coeffs_t) 117 | # out = self.layer4(out, coeffs_t) 118 | for module in self.layer1: 119 | out = module(out,coeffs_t) 120 | for module in self.layer2: 121 | out = module(out,coeffs_t) 122 | for module in self.layer3: 123 | out = module(out,coeffs_t) 124 | for module in self.layer4: 125 | out = module(out,coeffs_t) 126 | 127 | out = F.avg_pool2d(out, 4) 128 | out = out.view(out.size(0), -1) 129 | out = self.linear(out, coeffs_t) 130 | return out 131 | 132 | 133 | def ResNetCurve(num, fix_points, num_classes=10, initialize=False): 134 | if num == 18: 135 | return _ResNetCurve(BasicBlockCurves, [2,2,2,2], fix_points, num_classes, initialize) 136 | elif num == 34: 137 | return _ResNetCurve(BasicBlockCurves, [3,4,6,3], fix_points, num_classes, initialize) 138 | elif num == 50: 139 | return _ResNetCurve(BottleneckCurve, [3,4,6,3], fix_points, num_classes, initialize) 140 | elif num == 101: 141 | return _ResNetCurve(BottleneckCurve, [3,4,23,3], fix_points, num_classes, initialize) 142 | elif num == 152: 143 | return _ResNetCurve(BottleneckCurve, [3,8,36,3], fix_points, num_classes, initialize) 144 | else: 145 | raise NotImplementedError -------------------------------------------------------------------------------- /core/models/unet.py: -------------------------------------------------------------------------------- 1 | # https://github.com/milesial/Pytorch-UNet/blob/master/unet 2 | # https://blog.csdn.net/dongjinkun/article/details/115460931 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | class DoubleConv(nn.Module): 9 | """(convolution => [BN] => ReLU) * 2""" 10 | 11 | def __init__(self, in_channels, out_channels, mid_channels=None): 12 | super().__init__() 13 | if not mid_channels: 14 | mid_channels = out_channels 15 | self.double_conv = nn.Sequential( 16 | nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1), 17 | nn.BatchNorm2d(mid_channels), 18 | nn.ReLU(inplace=True), 19 | nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1), 20 | nn.BatchNorm2d(out_channels), 21 | nn.ReLU(inplace=True) 22 | ) 23 | 24 | def forward(self, x): 25 | return self.double_conv(x) 26 | 27 | class Down(nn.Module): 28 | """Downscaling with maxpool then double conv""" 29 | 30 | def __init__(self, in_channels, out_channels): 31 | super().__init__() 32 | self.maxpool_conv = nn.Sequential( 33 | nn.MaxPool2d(2), 34 | DoubleConv(in_channels, out_channels) 35 | ) 36 | 37 | def forward(self, x): 38 | return self.maxpool_conv(x) 39 | 40 | class Up(nn.Module): 41 | """Upscaling then double conv""" 42 | 43 | def __init__(self, in_channels, out_channels, bilinear=True): 44 | super().__init__() 45 | 46 | # if bilinear, use the normal convolutions to reduce the number of channels 47 | if bilinear: 48 | self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) # 双线性插值 49 | self.conv = DoubleConv(in_channels, out_channels, in_channels // 2) 50 | else: 51 | self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) # 转置卷积 52 | self.conv = DoubleConv(in_channels, out_channels) 53 | 54 | def forward(self, x1, x2): 55 | x1 = self.up(x1) 56 | # input is CHW 57 | diffY = x2.size()[2] - x1.size()[2] 58 | diffX = x2.size()[3] - x1.size()[3] 59 | 60 | x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, 61 | diffY // 2, diffY - diffY // 2]) 62 | # if you have padding issues, see 63 | # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a 64 | # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd 65 | x = torch.cat([x2, x1], dim=1) 66 | # x = torch.cat([x1, x1], dim=1) 67 | return self.conv(x) 68 | 69 | class OutConv(nn.Module): 70 | def __init__(self, in_channels, out_channels): 71 | super(OutConv, self).__init__() 72 | self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) 73 | # self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) 74 | 75 | def forward(self, x): 76 | return self.conv(x) 77 | 78 | class UNet(nn.Module): 79 | def __init__(self, args, n_channels, n_classes, bilinear=True): 80 | super(UNet, self).__init__() 81 | self.n_channels = n_channels 82 | self.n_classes = n_classes 83 | self.bilinear = bilinear 84 | 85 | """DoubleConv <-> (convolution => [BN] => ReLU) * 2""" 86 | self.inc = DoubleConv(n_channels, 64) 87 | self.down1 = Down(64, 128) 88 | self.down2 = Down(128, 256) 89 | self.down3 = Down(256, 512) 90 | factor = 2 if bilinear else 1 91 | self.down4 = Down(512, 1024 // factor) 92 | self.up1 = Up(1024, 512 // factor, bilinear) 93 | self.up2 = Up(512, 256 // factor, bilinear) 94 | self.up3 = Up(256, 128 // factor, bilinear) 95 | self.up4 = Up(128, 64, bilinear) 96 | self.outc = OutConv(64, n_classes) 97 | 98 | def forward(self, x): 99 | x1 = self.inc(x) 100 | x2 = self.down1(x1) 101 | x3 = self.down2(x2) 102 | x4 = self.down3(x3) 103 | x5 = self.down4(x4) 104 | x = self.up1(x5, x4) 105 | x = self.up2(x, x3) 106 | x = self.up3(x, x2) 107 | x = self.up4(x, x1) 108 | logits = self.outc(x) 109 | return logits 110 | 111 | class UNetLittle(nn.Module): 112 | def __init__(self, args, n_channels, n_classes, bilinear=True, first_channels=32): 113 | super(UNetLittle, self).__init__() 114 | self.n_channels = n_channels 115 | self.n_classes = n_classes 116 | self.bilinear = bilinear 117 | 118 | """DoubleConv <-> (convolution => [BN] => ReLU) * 2""" 119 | self.inc = DoubleConv(n_channels, first_channels) 120 | self.down1 = Down(first_channels, first_channels*2) 121 | self.down2 = Down(first_channels*2, first_channels*4) 122 | self.down3 = Down(first_channels*4, first_channels*8) 123 | factor = 2 if bilinear else 1 124 | self.down4 = Down(first_channels*8, first_channels*16 // factor) 125 | self.up1 = Up(first_channels*16, first_channels*8 // factor, bilinear) 126 | self.up2 = Up(first_channels*8, first_channels*4 // factor, bilinear) 127 | self.up3 = Up(first_channels*4, first_channels*2 // factor, bilinear) 128 | self.up4 = Up(first_channels*2, first_channels, bilinear) 129 | self.outc = OutConv(first_channels, n_classes) 130 | 131 | def forward(self, x): 132 | x1 = self.inc(x) 133 | x2 = self.down1(x1) 134 | x3 = self.down2(x2) 135 | x4 = self.down3(x3) 136 | x5 = self.down4(x4) 137 | x = self.up1(x5, x4) 138 | x = self.up2(x, x3) 139 | x = self.up3(x, x2) 140 | x = self.up4(x, x1) 141 | logits = self.outc(x) 142 | return logits 143 | 144 | class UNet3Layer(nn.Module): 145 | def __init__(self, args, n_channels, n_classes, bilinear=True): 146 | super(UNet3Layer, self).__init__() 147 | self.n_channels = n_channels 148 | self.n_classes = n_classes 149 | self.bilinear = bilinear 150 | 151 | """DoubleConv <-> (convolution => [BN] => ReLU) * 2""" 152 | self.inc = DoubleConv(n_channels, 64) 153 | self.down1 = Down(64, 128) 154 | self.down2 = Down(128, 256) 155 | factor = 2 if bilinear else 1 156 | self.down3 = Down(256, 512 // factor) 157 | self.up1 = Up(512, 256 // factor, bilinear) 158 | self.up2 = Up(256, 128 // factor, bilinear) 159 | self.up3 = Up(128, 64, bilinear) 160 | self.outc = OutConv(64, n_classes) 161 | 162 | def forward(self, x): 163 | x1 = self.inc(x) 164 | x2 = self.down1(x1) 165 | x3 = self.down2(x2) 166 | x4 = self.down3(x3) 167 | x = self.up1(x4, x3) 168 | x = self.up2(x, x2) 169 | x = self.up3(x, x1) 170 | logits = self.outc(x) 171 | return logits 172 | 173 | class UNet5Layer(nn.Module): 174 | def __init__(self, args, n_channels, n_classes, bilinear=True): 175 | super(UNet5Layer, self).__init__() 176 | self.n_channels = n_channels 177 | self.n_classes = n_classes 178 | self.bilinear = bilinear 179 | 180 | """DoubleConv <-> (convolution => [BN] => ReLU) * 2""" 181 | self.inc = DoubleConv(n_channels, 64) 182 | self.down1 = Down(64, 128) 183 | self.down2 = Down(128, 256) 184 | self.down3 = Down(256, 512) 185 | self.down4 = Down(512, 1024) 186 | factor = 2 if bilinear else 1 187 | self.down5 = Down(1024, 2048 // factor) 188 | self.up1 = Up(2048, 1024 // factor, bilinear) 189 | self.up2 = Up(1024, 512 // factor, bilinear) 190 | self.up3 = Up(512, 256 // factor, bilinear) 191 | self.up4 = Up(256, 128 // factor, bilinear) 192 | self.up5 = Up(128, 64, bilinear) 193 | self.outc = OutConv(64, n_classes) 194 | 195 | def forward(self, x): 196 | x1 = self.inc(x) 197 | x2 = self.down1(x1) 198 | x3 = self.down2(x2) 199 | x4 = self.down3(x3) 200 | x5 = self.down4(x4) 201 | x6 = self.down5(x5) 202 | x = self.up1(x6, x5) 203 | x = self.up2(x, x4) 204 | x = self.up3(x, x3) 205 | x = self.up4(x, x2) 206 | x = self.up5(x, x1) 207 | logits = self.outc(x) 208 | return logits 209 | 210 | if __name__ == '__main__': 211 | x = torch.randn(size=(128,3,224,224)).cuda() 212 | 213 | unet = UNetLittle(args=None, n_channels=3, n_classes=3, first_channels=16).cuda() 214 | 215 | 216 | y = unet(x) 217 | 218 | while True: 219 | pass -------------------------------------------------------------------------------- /core/models/vgg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import math 5 | 6 | 7 | class VGG(nn.Module): 8 | 9 | def __init__(self, features, num_classes=43): 10 | super(VGG, self).__init__() 11 | self.features = features 12 | self.classifier = nn.Linear(512, num_classes) 13 | self._initialize_weights() 14 | 15 | def forward(self, x): 16 | x = self.features(x) 17 | x = x.view(x.size(0), -1) 18 | x = self.classifier(x) 19 | return x 20 | 21 | def _initialize_weights(self): 22 | for m in self.modules(): 23 | if isinstance(m, nn.Conv2d): 24 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 25 | m.weight.data.normal_(0, math.sqrt(2. / n)) 26 | if m.bias is not None: 27 | m.bias.data.zero_() 28 | elif isinstance(m, nn.BatchNorm2d): 29 | m.weight.data.fill_(1) 30 | m.bias.data.zero_() 31 | elif isinstance(m, nn.Linear): 32 | n = m.weight.size(1) 33 | m.weight.data.normal_(0, 0.01) 34 | m.bias.data.zero_() 35 | 36 | 37 | def make_layers(cfg, batch_norm=False): 38 | layers = [] 39 | in_channels = 3 40 | for v in cfg: 41 | if v == 'M': 42 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 43 | else: 44 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 45 | if batch_norm: 46 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 47 | else: 48 | layers += [conv2d, nn.ReLU(inplace=True)] 49 | in_channels = v 50 | return nn.Sequential(*layers) 51 | 52 | 53 | cfg = { 54 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 55 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 56 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 57 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 58 | } 59 | 60 | 61 | def vgg11(num_classes=10, **kwargs): 62 | """VGG 11-layer model (configuration "A") 63 | Args: 64 | pretrained (bool): If True, returns a model pre-trained on ImageNet 65 | """ 66 | model = VGG(make_layers(cfg['A']), num_classes=num_classes, **kwargs) 67 | return model 68 | 69 | 70 | def vgg11_bn(num_classes=10, **kwargs): 71 | """VGG 11-layer model (configuration "A") with batch normalization""" 72 | model = VGG(make_layers(cfg['A'], batch_norm=True), num_classes=num_classes, **kwargs) 73 | return model 74 | 75 | 76 | def vgg13(num_classes=10, **kwargs): 77 | """VGG 13-layer model (configuration "B") 78 | Args: 79 | pretrained (bool): If True, returns a model pre-trained on ImageNet 80 | """ 81 | model = VGG(make_layers(cfg['B']), num_classes=num_classes, **kwargs) 82 | return model 83 | 84 | 85 | def vgg13_bn(num_classes=10, **kwargs): 86 | """VGG 13-layer model (configuration "B") with batch normalization""" 87 | model = VGG(make_layers(cfg['B'], batch_norm=True), num_classes=num_classes, **kwargs) 88 | return model 89 | 90 | 91 | def vgg16(num_classes=10, **kwargs): 92 | """VGG 16-layer model (configuration "D") 93 | Args: 94 | pretrained (bool): If True, returns a model pre-trained on ImageNet 95 | """ 96 | model = VGG(make_layers(cfg['D']), num_classes=num_classes, **kwargs) 97 | return model 98 | 99 | 100 | def vgg16_bn(num_classes=10, **kwargs): 101 | """VGG 16-layer model (configuration "D") with batch normalization""" 102 | model = VGG(make_layers(cfg['D'], batch_norm=True), num_classes=num_classes, **kwargs) 103 | return model 104 | 105 | 106 | def vgg19(num_classes=10, **kwargs): 107 | """VGG 19-layer model (configuration "E") 108 | Args: 109 | pretrained (bool): If True, returns a model pre-trained on ImageNet 110 | """ 111 | model = VGG(make_layers(cfg['E']), num_classes=num_classes, **kwargs) 112 | return model 113 | 114 | 115 | def vgg19_bn(num_classes=10, **kwargs): 116 | """VGG 19-layer model (configuration 'E') with batch normalization""" 117 | model = VGG(make_layers(cfg['E'], batch_norm=True), num_classes=num_classes, **kwargs) 118 | return model -------------------------------------------------------------------------------- /core/models/vgg_curve.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the part of the implement of model-repairing-based backdoor defense with MCR proposed in [1]. 3 | 4 | Reference: 5 | [1] Bridging Mode Connectivity in Loss Landscapes and Adversarial Robustness. ICLR, 2020. 6 | ''' 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import math 12 | import curves 13 | 14 | class VGGCurve(nn.Module): 15 | 16 | def __init__(self, features, num_classes=43, fix_points=None, initialize=False): 17 | super(VGGCurve, self).__init__() 18 | self.features = features 19 | self.classifier = curves.Linear(512, num_classes, fix_points) 20 | # self._initialize_weights() 21 | if initialize: 22 | for m in self.modules(): 23 | if isinstance(m, curves.Conv2d): 24 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 25 | for i in range(m.num_bends): 26 | getattr(m, 'weight_%d' % i).data.normal_(0, math.sqrt(2. / n)) 27 | getattr(m, 'bias_%d' % i).data.zero_() 28 | 29 | 30 | def forward(self, x, coeffs_t): 31 | # x = self.features(x) 32 | for module in self.features: 33 | if isinstance(module, curves.CurveModule): 34 | x = module(x, coeffs_t) 35 | else: 36 | x = module(x) 37 | x = x.view(x.size(0), -1) 38 | x = self.classifier(x, coeffs_t) 39 | return x 40 | 41 | #TODO: check if this is needed 42 | def _initialize_weights(self): 43 | for m in self.modules(): 44 | if isinstance(m, curves.Conv2d): 45 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 46 | m.weight.data.normal_(0, math.sqrt(2. / n)) 47 | if m.bias is not None: 48 | m.bias.data.zero_() 49 | elif isinstance(m, curves.BatchNorm2d): 50 | m.weight.data.fill_(1) 51 | m.bias.data.zero_() 52 | elif isinstance(m, curves.Linear): 53 | n = m.weight.size(1) 54 | m.weight.data.normal_(0, 0.01) 55 | m.bias.data.zero_() 56 | 57 | 58 | def make_layers(cfg, fix_points=None, batch_norm=False): 59 | layers = [] 60 | in_channels = 3 61 | for v in cfg: 62 | if v == 'M': 63 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 64 | else: 65 | conv2d = curves.Conv2d(in_channels, v, kernel_size=3, fix_points=fix_points, padding=1) 66 | if batch_norm: 67 | layers += [conv2d, curves.BatchNorm2d(v, fix_points), nn.ReLU(inplace=True)] 68 | else: 69 | layers += [conv2d, nn.ReLU(inplace=True)] 70 | in_channels = v 71 | return nn.Sequential(*layers) 72 | 73 | 74 | cfg = { 75 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 76 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 77 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 78 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 79 | } 80 | 81 | 82 | def vgg11(num_classes=10, fix_points=None, **kwargs): 83 | """VGG 11-layer model (configuration "A") 84 | Args: 85 | pretrained (bool): If True, returns a model pre-trained on ImageNet 86 | """ 87 | model = VGGCurve(make_layers(cfg['A'], fix_points), num_classes=num_classes, fix_points=fix_points, **kwargs) 88 | return model 89 | 90 | 91 | def vgg11_bn(num_classes=10, fix_points=None, **kwargs): 92 | """VGG 11-layer model (configuration "A") with batch normalization""" 93 | model = VGGCurve(make_layers(cfg['A'], fix_points, batch_norm=True), num_classes=num_classes, fix_points=fix_points, **kwargs) 94 | return model 95 | 96 | 97 | def vgg13(num_classes=10, fix_points=None, **kwargs): 98 | """VGG 13-layer model (configuration "B") 99 | Args: 100 | pretrained (bool): If True, returns a model pre-trained on ImageNet 101 | """ 102 | model = VGGCurve(make_layers(cfg['B'], fix_points), num_classes=num_classes, fix_points=fix_points, **kwargs) 103 | return model 104 | 105 | 106 | def vgg13_bn(num_classes=10, fix_points=None, **kwargs): 107 | """VGG 13-layer model (configuration "B") with batch normalization""" 108 | model = VGGCurve(make_layers(cfg['B'], fix_points, batch_norm=True), num_classes=num_classes, fix_points=fix_points, **kwargs) 109 | return model 110 | 111 | 112 | def vgg16(num_classes=10, fix_points=None, **kwargs): 113 | """VGG 16-layer model (configuration "D") 114 | Args: 115 | pretrained (bool): If True, returns a model pre-trained on ImageNet 116 | """ 117 | model = VGGCurve(make_layers(cfg['D'], fix_points), num_classes=num_classes, fix_points=fix_points, **kwargs) 118 | return model 119 | 120 | 121 | def vgg16_bn(num_classes=10, fix_points=None, **kwargs): 122 | """VGG 16-layer model (configuration "D") with batch normalization""" 123 | model = VGGCurve(make_layers(cfg['D'], fix_points, batch_norm=True), num_classes=num_classes, fix_points=fix_points, **kwargs) 124 | return model 125 | 126 | 127 | def vgg19(num_classes=10, fix_points=None, **kwargs): 128 | """VGG 19-layer model (configuration "E") 129 | Args: 130 | pretrained (bool): If True, returns a model pre-trained on ImageNet 131 | """ 132 | model = VGGCurve(make_layers(cfg['E'], fix_points), num_classes=num_classes, fix_points=fix_points, **kwargs) 133 | return model 134 | 135 | 136 | def vgg19_bn(num_classes=10, fix_points=None, **kwargs): 137 | """VGG 19-layer model (configuration 'E') with batch normalization""" 138 | model = VGGCurve(make_layers(cfg['E'], fix_points, batch_norm=True), num_classes=num_classes, fix_points=fix_points, **kwargs) 139 | return model -------------------------------------------------------------------------------- /core/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .accuracy import accuracy 2 | from .any2tensor import any2tensor 3 | from .log import Log 4 | from .test import test 5 | from .torchattacks import PGD 6 | from .supconloss import SupConLoss 7 | 8 | __all__ = [ 9 | 'Log', 'PGD', 'any2tensor', 'test', 'accuracy', 'SupConLoss' 10 | ] -------------------------------------------------------------------------------- /core/utils/accuracy.py: -------------------------------------------------------------------------------- 1 | def accuracy(output, target, topk=(1,)): 2 | """Computes the precision@k for the specified values of k""" 3 | maxk = max(topk) 4 | batch_size = target.size(0) 5 | 6 | _, pred = output.topk(maxk, 1, True, True) 7 | pred = pred.t() 8 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 9 | 10 | res = [] 11 | for k in topk: 12 | correct_k = correct[:k].contiguous().view(-1).float().sum(0) 13 | res.append(correct_k.mul_(100.0 / batch_size)) 14 | return res -------------------------------------------------------------------------------- /core/utils/any2tensor.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import PIL 3 | import numpy 4 | import torch 5 | from torchvision.transforms import functional as F 6 | 7 | 8 | def _any2tensor(x): 9 | """Convert a strpath, PIL.Image.Image, numpy.ndarray, torch.Tensor object to a torch.Tensor object. 10 | 11 | Args: 12 | x (strpath | PIL.Image.Image | numpy.ndarray | torch.Tensor): numpy.ndarray and torch.Tensor can have any shape. 13 | Hint: For strpath, x is converted to a torch.Tensor with shape (C, H, W), the channel order is decided by opencv. 14 | For PIL.Image.Image, x is converted to a torch.Tensor with shape (C, H, W), the channel order is decided by x itself. 15 | The channel order between opencv and PIL is different. 16 | 17 | Returns: 18 | torch.Tensor: The converted object. 19 | """ 20 | if type(x) == str: 21 | tmp = cv2.imread(x, cv2.IMREAD_UNCHANGED) 22 | if tmp.ndim == 2: 23 | return torch.from_numpy(tmp.reshape(1, tmp.shape[0], tmp.shape[1])) 24 | else: 25 | return torch.from_numpy(tmp.transpose((2, 0, 1))) 26 | elif type(x) == PIL.Image.Image: 27 | return F.pil_to_tensor(x) 28 | elif type(x) == numpy.ndarray: 29 | return torch.from_numpy(x) 30 | elif type(x) == torch.Tensor: 31 | return x.clone().detach() 32 | else: 33 | raise TypeError('x is an unsupported type, x should be strpath or PIL.Image.Image or numpy.ndarray or torch.Tensor. But got {}'.format(type(x))) 34 | 35 | 36 | def any2tensor(imgs): 37 | """Convert strpath, PIL.Image.Image, numpy.ndarray, torch.Tensor image(s) to a torch.Tensor. 38 | 39 | Args: 40 | imgs (list[strpath] | list[PIL.Image.Image] | list[numpy.ndarray] | list[torch.Tensor] | strpath | PIL.Image.Image | numpy.ndarray | torch.Tensor): The input images. 41 | 42 | Returns: 43 | torch.Tensor: The converted image(s). 44 | """ 45 | if isinstance(imgs, list): 46 | return torch.stack([_any2tensor(img) for img in imgs], dim=0) 47 | elif isinstance(imgs, (str, PIL.Image.Image, numpy.ndarray, torch.Tensor)): 48 | return _any2tensor(imgs) 49 | else: 50 | raise TypeError('imgs is an unsupported type, imgs should be list[strpath] | list[PIL.Image.Image] | list[numpy.ndarray] | list[torch.Tensor] | strpath | PIL.Image.Image | numpy.ndarray | torch.Tensor. But got {}'.format(type(imgs))) 51 | -------------------------------------------------------------------------------- /core/utils/compute_metric.py: -------------------------------------------------------------------------------- 1 | #-*- coding:utf-8 –*- 2 | import numpy as np 3 | 4 | # 计算混淆矩阵 5 | def compute_confusion_matrix(precited,expected): 6 | predicted = np.array(precited,dtype = int) 7 | expected = np.array(expected,dtype = int) 8 | part = precited**expected # 对结果进行分类，亦或使得判断正确的为0,判断错误的为1 9 | pcount = np.bincount(part) # 分类结果统计，pcount[0]为0的个数，pcount[1]为1的个数 10 | tp_list = list(precited & expected) # 将TP的计算结果转换为list 11 | fp_list = list(precited & ~expected) # 将FP的计算结果转换为list 12 | tp = tp_list.count(1) # 统计TP的个数 13 | fp = fp_list.count(1) # 统计FP的个数 14 | tn = pcount[0] - tp # 统计TN的个数 15 | fn = pcount[1] - fp # 统计FN的个数 16 | return tp, fp, tn, fn 17 | 18 | # 计算常用指标 19 | def compute_indexes(tp, fp, tn, fn): 20 | accuracy = (tp+tn) / (tp+tn+fp+fn) # 准确率 21 | precision = tp / (tp+fp) # 精确率 22 | recall = tp / (tp+fn) # 召回率 23 | F1 = (2*precision*recall) / (precision+recall) # F1 24 | return accuracy, precision, recall, F1 25 | -------------------------------------------------------------------------------- /core/utils/log.py: -------------------------------------------------------------------------------- 1 | class Log: 2 | def __init__(self, log_path): 3 | self.log_path = log_path 4 | 5 | def __call__(self, msg): 6 | print(msg, end='\n') 7 | with open(self.log_path,'a') as f: 8 | f.write(msg) 9 | -------------------------------------------------------------------------------- /core/utils/supconloss.py: -------------------------------------------------------------------------------- 1 | """ 2 | Author: Yonglong Tian (yonglong@mit.edu) 3 | Date: May 07, 2020 4 | """ 5 | from __future__ import print_function 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | 11 | class SupConLoss(nn.Module): 12 | """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. 13 | It also supports the unsupervised contrastive loss in SimCLR""" 14 | def __init__(self, temperature=0.07, contrast_mode='all', 15 | base_temperature=0.07): 16 | super(SupConLoss, self).__init__() 17 | self.temperature = temperature 18 | self.contrast_mode = contrast_mode 19 | self.base_temperature = base_temperature 20 | 21 | def forward(self, features, labels=None, mask=None): 22 | """Compute loss for model. If both `labels` and `mask` are None, 23 | it degenerates to SimCLR unsupervised loss: 24 | https://arxiv.org/pdf/2002.05709.pdf 25 | 26 | Args: 27 | features: hidden vector of shape [bsz, n_views, ...]. 28 | labels: ground truth of shape [bsz]. 29 | mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j 30 | has the same class as sample i. Can be asymmetric. 31 | Returns: 32 | A loss scalar. 33 | """ 34 | device = (torch.device('cuda') 35 | if features.is_cuda 36 | else torch.device('cpu')) 37 | 38 | if len(features.shape) < 3: 39 | raise ValueError('`features` needs to be [bsz, n_views, ...],' 40 | 'at least 3 dimensions are required') 41 | if len(features.shape) > 3: 42 | features = features.view(features.shape[0], features.shape[1], -1) 43 | 44 | batch_size = features.shape[0] 45 | if labels is not None and mask is not None: 46 | raise ValueError('Cannot define both `labels` and `mask`') 47 | elif labels is None and mask is None: 48 | mask = torch.eye(batch_size, dtype=torch.float32).to(device) 49 | elif labels is not None: 50 | labels = labels.contiguous().view(-1, 1) 51 | if labels.shape[0] != batch_size: 52 | raise ValueError('Num of labels does not match num of features') 53 | mask = torch.eq(labels, labels.T).float().to(device) 54 | else: 55 | mask = mask.float().to(device) 56 | 57 | contrast_count = features.shape[1] 58 | contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0) 59 | if self.contrast_mode == 'one': 60 | anchor_feature = features[:, 0] 61 | anchor_count = 1 62 | elif self.contrast_mode == 'all': 63 | anchor_feature = contrast_feature 64 | anchor_count = contrast_count 65 | else: 66 | raise ValueError('Unknown mode: {}'.format(self.contrast_mode)) 67 | 68 | # compute logits 69 | anchor_dot_contrast = torch.div( 70 | torch.matmul(anchor_feature, contrast_feature.T), 71 | self.temperature) 72 | # for numerical stability 73 | logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True) 74 | logits = anchor_dot_contrast - logits_max.detach() 75 | 76 | # tile mask 77 | mask = mask.repeat(anchor_count, contrast_count) 78 | # mask-out self-contrast cases 79 | logits_mask = torch.scatter( 80 | torch.ones_like(mask), 81 | 1, 82 | torch.arange(batch_size * anchor_count).view(-1, 1).to(device), 83 | 0 84 | ) 85 | mask = mask * logits_mask 86 | 87 | # compute log_prob 88 | exp_logits = torch.exp(logits) * logits_mask 89 | log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True)) 90 | 91 | # compute mean of log-likelihood over positive 92 | # modified to handle edge cases when there is no positive pair 93 | # for an anchor point. 94 | # Edge case e.g.:- 95 | # features of shape: [4,1,...] 96 | # labels: [0,1,1,2] 97 | # loss before mean: [nan, ..., ..., nan] 98 | mask_pos_pairs = mask.sum(1) 99 | mask_pos_pairs = torch.where(mask_pos_pairs < 1e-6, 1, mask_pos_pairs) 100 | mean_log_prob_pos = (mask * log_prob).sum(1) / mask_pos_pairs 101 | 102 | # loss 103 | loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos 104 | loss = loss.view(anchor_count, batch_size).mean() 105 | 106 | return loss 107 | -------------------------------------------------------------------------------- /core/utils/test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import random 4 | import time 5 | 6 | import numpy as np 7 | import torch 8 | import torch.nn as nn 9 | from torch.utils.data import DataLoader 10 | from torchvision.datasets import CIFAR10, MNIST, DatasetFolder 11 | 12 | from .accuracy import accuracy 13 | from .log import Log 14 | 15 | 16 | def _seed_worker(worker_id): 17 | worker_seed = torch.initial_seed() % 2**32 18 | np.random.seed(worker_seed) 19 | random.seed(worker_seed) 20 | 21 | 22 | def _test(model, dataset, device, batch_size=16, num_workers=8): 23 | with torch.no_grad(): 24 | test_loader = DataLoader( 25 | dataset, 26 | batch_size=batch_size, 27 | shuffle=False, 28 | num_workers=num_workers, 29 | drop_last=False, 30 | pin_memory=True, 31 | worker_init_fn=_seed_worker 32 | ) 33 | 34 | model = model.to(device) 35 | model.eval() 36 | 37 | predict_digits = [] 38 | labels = [] 39 | for batch in test_loader: 40 | batch_img, batch_label = batch 41 | batch_img = batch_img.to(device) 42 | batch_img = model(batch_img) 43 | batch_img = batch_img.cpu() 44 | predict_digits.append(batch_img) 45 | labels.append(batch_label) 46 | 47 | predict_digits = torch.cat(predict_digits, dim=0) 48 | labels = torch.cat(labels, dim=0) 49 | return predict_digits, labels 50 | 51 | 52 | def test(model, dataset, schedule): 53 | """Uniform test API for any model and any dataset. 54 | 55 | Args: 56 | model (torch.nn.Module): Network. 57 | dataset (torch.utils.data.Dataset): Dataset. 58 | schedule (dict): Testing schedule. 59 | """ 60 | 61 | if 'test_model' in schedule: 62 | model.load_state_dict(torch.load(schedule['test_model']), strict=False) 63 | 64 | work_dir = osp.join(schedule['save_dir'], schedule['experiment_name'] + '_' + time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())) 65 | os.makedirs(work_dir, exist_ok=True) 66 | log = Log(osp.join(work_dir, 'log.txt')) 67 | 68 | # Use GPU 69 | if 'device' in schedule and schedule['device'] == 'GPU': 70 | log('==========Use GPUs to train==========\n') 71 | 72 | CUDA_VISIBLE_DEVICES = '' 73 | if 'CUDA_VISIBLE_DEVICES' in os.environ: 74 | CUDA_VISIBLE_DEVICES = os.environ['CUDA_VISIBLE_DEVICES'] 75 | else: 76 | CUDA_VISIBLE_DEVICES = ','.join([str(i) for i in range(torch.cuda.device_count())]) 77 | log(f'CUDA_VISIBLE_DEVICES={CUDA_VISIBLE_DEVICES}\n') 78 | 79 | if CUDA_VISIBLE_DEVICES == '': 80 | raise ValueError(f'This machine has no visible cuda devices!') 81 | 82 | CUDA_SELECTED_DEVICES = '' 83 | if 'CUDA_SELECTED_DEVICES' in schedule: 84 | CUDA_SELECTED_DEVICES = schedule['CUDA_SELECTED_DEVICES'] 85 | else: 86 | CUDA_SELECTED_DEVICES = CUDA_VISIBLE_DEVICES 87 | log(f'CUDA_SELECTED_DEVICES={CUDA_SELECTED_DEVICES}\n') 88 | 89 | CUDA_VISIBLE_DEVICES_LIST = sorted(CUDA_VISIBLE_DEVICES.split(',')) 90 | CUDA_SELECTED_DEVICES_LIST = sorted(CUDA_SELECTED_DEVICES.split(',')) 91 | 92 | CUDA_VISIBLE_DEVICES_SET = set(CUDA_VISIBLE_DEVICES_LIST) 93 | CUDA_SELECTED_DEVICES_SET = set(CUDA_SELECTED_DEVICES_LIST) 94 | if not (CUDA_SELECTED_DEVICES_SET <= CUDA_VISIBLE_DEVICES_SET): 95 | raise ValueError(f'CUDA_VISIBLE_DEVICES should be a subset of CUDA_VISIBLE_DEVICES!') 96 | 97 | GPU_num = len(CUDA_SELECTED_DEVICES_SET) 98 | device_ids = [CUDA_VISIBLE_DEVICES_LIST.index(CUDA_SELECTED_DEVICE) for CUDA_SELECTED_DEVICE in CUDA_SELECTED_DEVICES_LIST] 99 | device = torch.device(f'cuda:{device_ids[0]}') 100 | model = model.to(device) 101 | 102 | if GPU_num > 1: 103 | model = nn.DataParallel(model, device_ids=device_ids, output_device=device_ids[0]) 104 | # Use CPU 105 | else: 106 | device = torch.device("cpu") 107 | 108 | # device = torch.device("cuda:0") 109 | if schedule['metric'] == 'ASR_NoTarget': 110 | if isinstance(dataset, CIFAR10): 111 | data = [] 112 | targets = [] 113 | for i, target in enumerate(dataset.targets): 114 | if target != schedule['y_target']: 115 | data.append(dataset.data[i]) 116 | targets.append(target) 117 | data = np.stack(data, axis=0) 118 | dataset.data = data 119 | dataset.targets = targets 120 | elif isinstance(dataset, MNIST): 121 | data = [] 122 | targets = [] 123 | for i, target in enumerate(dataset.targets): 124 | if int(target) != schedule['y_target']: 125 | data.append(dataset.data[i]) 126 | targets.append(target) 127 | data = torch.stack(data, dim=0) 128 | dataset.data = data 129 | dataset.targets = targets 130 | elif isinstance(dataset, DatasetFolder): 131 | samples = [] 132 | for sample in dataset.samples: 133 | if sample[1] != schedule['y_target']: 134 | samples.append(sample) 135 | dataset.samples = samples 136 | else: 137 | raise NotImplementedError 138 | 139 | 140 | last_time = time.time() 141 | predict_digits, labels = _test(model, dataset, device, schedule['batch_size'], schedule['num_workers']) 142 | total_num = labels.size(0) 143 | prec1, prec5 = accuracy(predict_digits, labels, topk=(1, 5)) 144 | top1_correct = int(round(prec1.item() / 100.0 * total_num)) 145 | top5_correct = int(round(prec5.item() / 100.0 * total_num)) 146 | msg = f"==========Test result on {schedule['metric']}==========\n" + \ 147 | time.strftime("[%Y-%m-%d_%H:%M:%S] ", time.localtime()) + \ 148 | f"Top-1 correct / Total: {top1_correct}/{total_num}, Top-1 accuracy: {top1_correct/total_num}, Top-5 correct / Total: {top5_correct}/{total_num}, Top-5 accuracy: {top5_correct/total_num}, time: {time.time()-last_time}\n" 149 | log(msg) 150 | -------------------------------------------------------------------------------- /core/utils/torchattacks/README.md: -------------------------------------------------------------------------------- 1 | This directory contains PGD attack codes borrowed from [torchattacks](https://github.com/Harry24k/adversarial-attacks-pytorch). -------------------------------------------------------------------------------- /core/utils/torchattacks/__init__.py: -------------------------------------------------------------------------------- 1 | from .attacks.pgd import PGD 2 | 3 | __version__ = 2.6 4 | -------------------------------------------------------------------------------- /core/utils/torchattacks/attack.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class Attack(object): 5 | r""" 6 | Base class for all attacks. 7 | .. note:: 8 | It automatically set device to the device where given model is. 9 | It temporarily changes the original model's training mode to `test` 10 | by `.eval()` only during an attack process. 11 | """ 12 | 13 | def __init__(self, name, model): 14 | r""" 15 | Initializes internal attack state. 16 | Arguments: 17 | name (str) : name of an attack. 18 | model (torch.nn.Module): model to attack. 19 | """ 20 | 21 | self.attack = name 22 | self.model = model 23 | self.model_name = str(model).split("(")[0] 24 | 25 | self.training = model.training 26 | self.device = next(model.parameters()).device 27 | 28 | self._targeted = 1 29 | self._attack_mode = "original" 30 | self._return_type = "float" 31 | 32 | def forward(self, *input): 33 | r""" 34 | It defines the computation performed at every call. 35 | Should be overridden by all subclasses. 36 | """ 37 | raise NotImplementedError 38 | 39 | def set_attack_mode(self, mode): 40 | r""" 41 | Set the attack mode. 42 | 43 | Arguments: 44 | mode (str) : 'original' (DEFAULT) 45 | 'targeted' - Use input labels as targeted labels. 46 | 'least_likely' - Use least likely labels as targeted labels. 47 | """ 48 | if self._attack_mode == "only_original": 49 | raise ValueError( 50 | "Changing attack mode is not supported in this attack method." 51 | ) 52 | 53 | if mode == "original": 54 | self._attack_mode = "original" 55 | self._targeted = 1 56 | self._transform_label = self._get_label 57 | elif mode == "targeted": 58 | self._attack_mode = "targeted" 59 | self._targeted = -1 60 | self._transform_label = self._get_label 61 | elif mode == "least_likely": 62 | self._attack_mode = "least_likely" 63 | self._targeted = -1 64 | self._transform_label = self._get_least_likely_label 65 | else: 66 | raise ValueError( 67 | mode 68 | + " is not a valid mode. [Options : original, targeted, least_likely]" 69 | ) 70 | 71 | def set_return_type(self, type): 72 | r""" 73 | Set the return type of adversarial images: `int` or `float`. 74 | Arguments: 75 | type (str) : 'float' or 'int'. (DEFAULT : 'float') 76 | """ 77 | if type == "float": 78 | self._return_type = "float" 79 | elif type == "int": 80 | self._return_type = "int" 81 | else: 82 | raise ValueError(type + " is not a valid type. [Options : float, int]") 83 | 84 | def save(self, save_path, data_loader, verbose=True): 85 | r""" 86 | Save adversarial images as torch.tensor from given torch.utils.data.DataLoader. 87 | Arguments: 88 | save_path (str) : save_path. 89 | data_loader (torch.utils.data.DataLoader) : data loader. 90 | verbose (bool) : True for displaying detailed information. (DEFAULT : True) 91 | """ 92 | self.model.eval() 93 | 94 | image_list = [] 95 | label_list = [] 96 | 97 | correct = 0 98 | total = 0 99 | 100 | total_batch = len(data_loader) 101 | 102 | for step, (images, labels) in enumerate(data_loader): 103 | adv_images = self.__call__(images, labels) 104 | 105 | image_list.append(adv_images.cpu()) 106 | label_list.append(labels.cpu()) 107 | 108 | if self._return_type == "int": 109 | adv_images = adv_images.float() / 255 110 | 111 | if verbose: 112 | outputs = self.model(adv_images) 113 | _, predicted = torch.max(outputs.data, 1) 114 | total += labels.size(0) 115 | correct += (predicted == labels.to(self.device)).sum() 116 | 117 | acc = 100 * float(correct) / total 118 | print( 119 | "- Save Progress : %2.2f %% / Accuracy : %2.2f %%" 120 | % ((step + 1) / total_batch * 100, acc), 121 | end="\r", 122 | ) 123 | 124 | x = torch.cat(image_list, 0) 125 | y = torch.cat(label_list, 0) 126 | torch.save((x, y), save_path) 127 | print("\n- Save Complete!") 128 | 129 | self._switch_model() 130 | 131 | def _transform_label(self, images, labels): 132 | r""" 133 | Function for changing the attack mode. 134 | """ 135 | return labels 136 | 137 | def _get_label(self, images, labels): 138 | r""" 139 | Function for changing the attack mode. 140 | Return input labels. 141 | """ 142 | return labels 143 | 144 | def _get_least_likely_label(self, images, labels): 145 | r""" 146 | Function for changing the attack mode. 147 | Return least likely labels. 148 | """ 149 | outputs = self.model(images) 150 | _, labels = torch.min(outputs.data, 1) 151 | labels = labels.detach_() 152 | return labels 153 | 154 | def _to_uint(self, images): 155 | r""" 156 | Function for changing the return type. 157 | Return images as int. 158 | """ 159 | return (images * 255).type(torch.uint8) 160 | 161 | def _switch_model(self): 162 | r""" 163 | Function for changing the training mode of the model. 164 | """ 165 | if self.training: 166 | self.model.train() 167 | else: 168 | self.model.eval() 169 | 170 | def __str__(self): 171 | info = self.__dict__.copy() 172 | 173 | del_keys = ["model", "attack"] 174 | 175 | for key in info.keys(): 176 | if key[0] == "_": 177 | del_keys.append(key) 178 | 179 | for key in del_keys: 180 | del info[key] 181 | 182 | info["attack_mode"] = self._attack_mode 183 | if info["attack_mode"] == "only_original": 184 | info["attack_mode"] = "original" 185 | 186 | info["return_type"] = self._return_type 187 | 188 | return ( 189 | self.attack 190 | + "(" 191 | + ", ".join("{}={}".format(key, val) for key, val in info.items()) 192 | + ")" 193 | ) 194 | 195 | def __call__(self, *input, **kwargs): 196 | self.model.eval() 197 | images = self.forward(*input, **kwargs) 198 | self._switch_model() 199 | 200 | if self._return_type == "int": 201 | images = self._to_uint(images) 202 | 203 | return images 204 | -------------------------------------------------------------------------------- /core/utils/torchattacks/attacks/__init__.py: -------------------------------------------------------------------------------- 1 | from .pgd import PGD 2 | -------------------------------------------------------------------------------- /core/utils/torchattacks/attacks/pgd.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from ..attack import Attack 5 | 6 | 7 | class PGD(Attack): 8 | r""" 9 | PGD in the paper 'Towards Deep Learning Models Resistant to Adversarial Attacks' 10 | [https://arxiv.org/abs/1706.06083] 11 | 12 | Distance Measure : Linf 13 | Arguments: 14 | model (nn.Module): model to attack. 15 | eps (float): maximum perturbation. (DEFALUT : 0.3) 16 | alpha (float): step size. (DEFALUT : 2/255) 17 | steps (int): number of steps. (DEFALUT : 40) 18 | random_start (bool): using random initialization of delta. (DEFAULT : False) 19 | 20 | Shape: 21 | - images: :math:`(N, C, H, W)` where `N = number of batches`, `C = number of channels`, `H = height` and `W = width`. It must have a range [0, 1]. 22 | - labels: :math:`(N)` where each value :math:`y_i` is :math:`0 \leq y_i \leq` `number of labels`. 23 | - output: :math:`(N, C, H, W)`. 24 | 25 | Examples:: 26 | >>> attack = torchattacks.PGD(model, eps = 8/255, alpha = 1/255, steps=40, random_start=False) 27 | >>> adv_images = attack(images, labels) 28 | 29 | """ 30 | 31 | def __init__(self, model, eps=0.3, alpha=2 / 255, steps=40, random_start=False): 32 | super(PGD, self).__init__("PGD", model) 33 | self.eps = eps 34 | self.alpha = alpha 35 | self.steps = steps 36 | self.random_start = random_start 37 | 38 | def forward(self, images, labels): 39 | r""" 40 | Overridden. 41 | """ 42 | images = images.to(self.device) 43 | labels = labels.to(self.device) 44 | labels = self._transform_label(images, labels) 45 | loss = nn.CrossEntropyLoss() 46 | 47 | adv_images = images.clone().detach() 48 | 49 | if self.random_start: 50 | # Starting at a uniformly random point 51 | adv_images = adv_images + torch.empty_like(adv_images).uniform_( 52 | -self.eps, self.eps 53 | ) 54 | adv_images = torch.clamp(adv_images, min=0, max=1) 55 | 56 | for i in range(self.steps): 57 | adv_images.requires_grad = True 58 | outputs = self.model(adv_images) 59 | 60 | cost = self._targeted * loss(outputs, labels).to(self.device) 61 | 62 | grad = torch.autograd.grad( 63 | cost, adv_images, retain_graph=False, create_graph=False 64 | )[0] 65 | 66 | adv_images = adv_images.detach() + self.alpha * grad.sign() 67 | delta = torch.clamp(adv_images - images, min=-self.eps, max=self.eps) 68 | adv_images = torch.clamp(images + delta, min=0, max=1).detach() 69 | 70 | return adv_images 71 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | --find-links https://download.pytorch.org/whl/torch_stable.html 2 | tqdm==4.62.3 3 | matplotlib==3.5.1 4 | numpy>=1.19.5 5 | pillow>=10.2.0 6 | opencv-python>=4.5.1.48 7 | torch==1.8.0+cu111 8 | torchvision==0.9.0+cu111 9 | torchaudio==0.8.0 10 | scipy==1.7.3 11 | requests==2.26.0 12 | termcolor==1.1.0 13 | easydict==1.9.0 14 | seaborn==0.11.2 15 | imageio==2.16.0 16 | lpips==0.1.4 17 | scikit-learn==1.3.2 18 | -------------------------------------------------------------------------------- /tests/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUYimingLi/BackdoorBox/945ea2125096dd19720be1489d0d92239c0541fd/tests/.DS_Store -------------------------------------------------------------------------------- /tests/test_AdaptivePatch.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | sys.path.append(os.getcwd()) 4 | import cv2 5 | import torch 6 | import torch.nn as nn 7 | import torchvision 8 | from torch.utils.data import Dataset, dataloader 9 | import numpy as np 10 | from torchvision import transforms 11 | from torch.utils.data import DataLoader 12 | from torchvision.datasets import DatasetFolder 13 | import torchvision.models as models 14 | import core 15 | import argparse 16 | 17 | parser = argparse.ArgumentParser(description='PyTorch Attack') 18 | parser.add_argument('--gpu', default='0', type=str, choices=[str(x) for x in range(8)]) 19 | parser.add_argument('--dataset', default='CIFAR10', type=str, choices=['CIFAR10', 'ImageNet50']) 20 | parser.add_argument('--model', default='ResNet18', type=str, choices=['ResNet18', 'ResNet50', 'VGG16', 'InceptionV3', 'DenseNet121', 'ViT']) 21 | parser.add_argument('--attack', default='Adaptive', type=str, choices=['Benign', 'BadNets', 'Blended', 'WaNet', 'BATT', 'Physical', 'LC','Adaptive']) 22 | 23 | args = parser.parse_args() 24 | state = {k: v for k, v in args._get_kwargs()} 25 | 26 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 27 | 28 | # ========== Set global settings ========== 29 | global_seed = 666 30 | deterministic = True 31 | torch.manual_seed(global_seed) 32 | CUDA_VISIBLE_DEVICES = args.gpu 33 | dataset=args.dataset 34 | model=args.model 35 | attack=args.attack 36 | datasets_root_dir = f'./data/{dataset}' 37 | save_path = f'./{dataset}/{model}/{attack}' 38 | 39 | if dataset == 'CIFAR10': 40 | img_size = 32 41 | mean = [0.4914, 0.4822, 0.4465] 42 | std = [0.2023, 0.1994, 0.2010] 43 | num_classes = 10 44 | elif dataset == 'ImageNet50': 45 | img_size = 224 46 | mean = [0.485, 0.456, 0.406] 47 | std = [0.229, 0.224, 0.225] 48 | num_classes = 50 49 | else: 50 | raise NotImplementedError 51 | 52 | input_size = img_size 53 | 54 | if model == 'ResNet18': 55 | my_model = core.models.ResNet(18, num_classes=num_classes) 56 | if dataset == 'ImageNet50': 57 | my_model = models.resnet18(weights=None, num_classes=num_classes) 58 | lr = 0.1 59 | elif model == 'ResNet50': 60 | my_model = core.models.ResNet(50, num_classes=num_classes) 61 | if dataset == 'ImageNet50': 62 | my_model = models.resnet50(pretrained=False, num_classes=num_classes) 63 | lr = 0.1 64 | elif model == 'VGG16': 65 | deterministic = False 66 | my_model = core.models.vgg16(num_classes=num_classes) 67 | if dataset == 'ImageNet50': 68 | my_model = models.vgg16(pretrained=False, num_classes=num_classes) 69 | lr = 0.01 70 | elif model == 'InceptionV3': 71 | my_model = models.inception_v3(pretrained=False, num_classes=num_classes, aux_logits = False) # 299*299 72 | if dataset == 'CIFAR10': 73 | input_size = 96 74 | lr = 0.1 75 | elif model == 'DenseNet121': 76 | my_model = models.densenet121(pretrained=False, num_classes=num_classes) # 224*224 77 | lr = 0.1 78 | elif model == 'ViT': 79 | # my_model = models.vit_b_16(weights=None, num_classes=num_classes) # 224*224 80 | my_model = core.models.ViT( 81 | image_size = img_size, 82 | patch_size = int(img_size / 8), 83 | num_classes = num_classes, 84 | dim = int(512), 85 | depth = 6, 86 | heads = 8, 87 | mlp_dim = 512, 88 | dropout = 0.1, 89 | emb_dropout = 0.1 90 | ) 91 | lr = 1e-3 92 | # input_size = 224 93 | else: 94 | raise NotImplementedError 95 | 96 | my_model = my_model.to('cuda' if torch.cuda.is_available() else 'cpu') 97 | 98 | transform_train = transforms.Compose([ 99 | transforms.ToTensor(), 100 | transforms.RandomHorizontalFlip(), 101 | transforms.Resize(input_size), 102 | transforms.Normalize(mean, std), 103 | ]) 104 | 105 | trainset = DatasetFolder(root=os.path.join(datasets_root_dir, 'train'), 106 | transform=transform_train, 107 | loader=cv2.imread, 108 | extensions=('png','jpeg',), 109 | target_transform=None, 110 | is_valid_file=None, 111 | ) 112 | 113 | transform_test = transforms.Compose([ 114 | transforms.ToTensor(), 115 | transforms.Resize(input_size), 116 | transforms.Normalize(mean, std), 117 | ]) 118 | 119 | testset = DatasetFolder(root=os.path.join(datasets_root_dir, 'test'), 120 | transform=transform_test, 121 | loader=cv2.imread, 122 | extensions=('png','jpeg',), 123 | target_transform=None, 124 | is_valid_file=None, 125 | ) 126 | 127 | trigger_dir = './adaptive_triggers' 128 | trigger_names = [ 129 | f'phoenix_corner_{img_size}.png', 130 | f'badnet_patch4_{img_size}.png', 131 | f'firefox_corner_{img_size}.png', 132 | f'trojan_square_{img_size}.png', 133 | ] 134 | trigger_path = [os.path.join(trigger_dir, name) for name in trigger_names] 135 | 136 | patterns = [] 137 | for path in trigger_path: 138 | pattern = cv2.imread(path) 139 | pattern = transforms.ToTensor()(pattern) 140 | patterns.append(pattern) 141 | 142 | alphas = [0.5, 0.5, 0.2, 0.3] 143 | if model == 'ResNet18': 144 | poisoned_rate = 0.01 145 | covered_rate = 0.02 146 | elif model in ['VGG16', 'DenseNet121', 'ViT']: 147 | poisoned_rate = 0.03 148 | covered_rate = 0.06 149 | 150 | attacker = core.AdaptivePatch( 151 | train_dataset=trainset, 152 | test_dataset=testset, 153 | model=my_model, 154 | loss=nn.CrossEntropyLoss(), 155 | y_target=0, 156 | poisoned_rate=poisoned_rate, 157 | covered_rate=covered_rate, 158 | patterns=patterns, 159 | alphas=alphas, 160 | seed=global_seed, 161 | deterministic=deterministic, 162 | ) 163 | 164 | 165 | benign_training = False 166 | if attack == 'Benign': 167 | benign_training = True 168 | 169 | # Train Attacked Model 170 | schedule = { 171 | 'device': 'GPU', 172 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 173 | 'GPU_num': 1, 174 | 175 | 'benign_training': benign_training, 176 | 'batch_size': 128, 177 | 'num_workers': 8, 178 | 179 | 'lr': lr, 180 | 'momentum': 0.9, 181 | 'weight_decay': 5e-4, 182 | 'gamma': 0.1, 183 | 'schedule': [100, 130], 184 | 185 | 'epochs': 150, 186 | 187 | 'log_iteration_interval': 100, 188 | 'test_epoch_interval': 10, 189 | 'save_epoch_interval': 20, 190 | 191 | 'save_dir': save_path, 192 | 'experiment_name': f'Normalize_{model}_{dataset}_{attack}' 193 | } 194 | attacker.train(schedule) 195 | -------------------------------------------------------------------------------- /tests/test_BATT.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of benign training and poisoned training under BadNets. 3 | ''' 4 | 5 | 6 | import os 7 | 8 | import cv2 9 | import torch 10 | import torch.nn as nn 11 | from torch.utils.data import Dataset 12 | import torchvision 13 | from torchvision.transforms import Compose, ToTensor 14 | import torchvision.transforms as transforms 15 | 16 | import core 17 | 18 | global_seed = 666 19 | deterministic = True 20 | torch.manual_seed(global_seed) 21 | 22 | # ============== cifar10 ============== 23 | # Define Benign Training and Testing Dataset 24 | dataset = torchvision.datasets.CIFAR10 25 | 26 | transform_train = Compose([ 27 | ToTensor(), 28 | ]) 29 | trainset = dataset('data', train=True, transform=transform_train, download=True) 30 | 31 | transform_test = Compose([ 32 | ToTensor() 33 | ]) 34 | testset = dataset('data', train=False, transform=transform_test, download=True) 35 | 36 | batt = core.BATT( 37 | train_dataset=trainset, 38 | test_dataset=testset, 39 | model=core.models.ResNet(18), 40 | loss=nn.CrossEntropyLoss(), 41 | y_target=1, 42 | poisoned_rate=0.05, 43 | seed=global_seed, 44 | deterministic=deterministic 45 | ) 46 | 47 | poisoned_train_dataset, poisoned_test_dataset = batt.get_poisoned_dataset() 48 | 49 | # Train Infected Model 50 | schedule = { 51 | 'device': 'GPU', 52 | 'CUDA_VISIBLE_DEVICES': '0', 53 | 'GPU_num': 1, 54 | 55 | 'benign_training': False, # Train Infected Model 56 | 'batch_size': 128, 57 | 'num_workers': 8, 58 | 59 | 'lr': 0.1, 60 | 'momentum': 0.9, 61 | 'weight_decay': 5e-4, 62 | 'gamma': 0.1, 63 | 'schedule': [150, 180], 64 | 65 | 'epochs': 200, 66 | 67 | 'log_iteration_interval': 100, 68 | 'test_epoch_interval': 10, 69 | 'save_epoch_interval': 10, 70 | 71 | 'save_dir': './result/', 72 | 'experiment_name': 'batt_cifar10' 73 | } 74 | 75 | batt.train(schedule) 76 | infected_model = batt.get_model() 77 | 78 | 79 | # ============== mnist ============== 80 | # Define Benign Training and Testing Dataset 81 | dataset = torchvision.datasets.MNIST 82 | 83 | 84 | transform_train = Compose([ 85 | ToTensor(), 86 | ]) 87 | trainset = dataset('data', train=True, transform=transform_train, download=True) 88 | 89 | transform_test = Compose([ 90 | ToTensor() 91 | ]) 92 | testset = dataset('data', train=False, transform=transform_test, download=True) 93 | 94 | batt = core.BATT( 95 | train_dataset=trainset, 96 | test_dataset=testset, 97 | model=core.models.BaselineMNISTNetwork(), 98 | loss=nn.CrossEntropyLoss(), 99 | y_target=1, 100 | poisoned_rate=0.05, 101 | seed=global_seed, 102 | deterministic=deterministic 103 | ) 104 | 105 | poisoned_train_dataset, poisoned_test_dataset = batt.get_poisoned_dataset() 106 | 107 | # Train Infected Model 108 | schedule = { 109 | 'device': 'GPU', 110 | 'CUDA_VISIBLE_DEVICES': '0', 111 | 'GPU_num': 1, 112 | 113 | 'benign_training': False, # Train Infected Model 114 | 'batch_size': 128, 115 | 'num_workers': 8, 116 | 117 | 'lr': 0.1, 118 | 'momentum': 0.9, 119 | 'weight_decay': 5e-4, 120 | 'gamma': 0.1, 121 | 'schedule': [150, 180], 122 | 123 | 'epochs': 200, 124 | 125 | 'log_iteration_interval': 100, 126 | 'test_epoch_interval': 10, 127 | 'save_epoch_interval': 10, 128 | 129 | 'save_dir': './result/', 130 | 'experiment_name': 'batt_mnist' 131 | } 132 | 133 | batt.train(schedule) 134 | infected_model = batt.get_model() 135 | 136 | 137 | 138 | # ============== GTSRB ============== 139 | # Define Benign Training and Testing Dataset 140 | dataset = torchvision.datasets.DatasetFolder 141 | 142 | transform_train = Compose([ 143 | transforms.ToPILImage(), 144 | transforms.Resize((32,32)), 145 | ToTensor() 146 | ]) 147 | trainset = dataset( 148 | root='/data1/xutong/dataset/GTSRB/Train', 149 | loader=cv2.imread, 150 | extensions=('png',), 151 | transform=transform_train, 152 | target_transform=None, 153 | is_valid_file=None) 154 | 155 | transform_test = Compose([ 156 | transforms.ToPILImage(), 157 | transforms.Resize((32,32)), 158 | ToTensor() 159 | ]) 160 | testset = dataset( 161 | root='/data1/xutong/dataset/GTSRB/testset', 162 | loader=cv2.imread, 163 | extensions=('png',), 164 | transform=transform_test, 165 | target_transform=None, 166 | is_valid_file=None) 167 | 168 | batt = core.BATT( 169 | train_dataset=trainset, 170 | test_dataset=testset, 171 | model=core.models.ResNet(18,43), 172 | loss=nn.CrossEntropyLoss(), 173 | y_target=1, 174 | poisoned_rate=0.05, 175 | poisoned_transform_train_index=2, 176 | poisoned_transform_test_index=2, 177 | schedule=None, 178 | seed=666, 179 | ) 180 | 181 | poisoned_train_dataset, poisoned_test_dataset = batt.get_poisoned_dataset() 182 | 183 | # Train Infected Model 184 | schedule = { 185 | 'device': 'GPU', 186 | 'CUDA_VISIBLE_DEVICES': '0', 187 | 'GPU_num': 1, 188 | 189 | 'benign_training': False, # Train Infected Model 190 | 'batch_size': 64, 191 | 'num_workers': 8, 192 | 193 | 'lr': 0.1, 194 | 'momentum': 0.9, 195 | 'weight_decay': 5e-4, 196 | 'gamma': 0.1, 197 | 'schedule': [150, 180], 198 | 199 | 'epochs': 30, 200 | 201 | 'log_iteration_interval': 100, 202 | 'test_epoch_interval': 10, 203 | 'save_epoch_interval': 10, 204 | 205 | 'save_dir': './result/', 206 | 'experiment_name': 'batt_gtsrb' 207 | } 208 | 209 | batt.train(schedule) 210 | infected_model = batt.get_model() 211 | -------------------------------------------------------------------------------- /tests/test_BadNets.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of poisoned training under BadNets. 3 | ''' 4 | 5 | 6 | import os.path as osp 7 | 8 | import cv2 9 | import torch 10 | import torch.nn as nn 11 | import torchvision 12 | from torchvision.datasets import DatasetFolder 13 | from torchvision.transforms import Compose, ToTensor, RandomHorizontalFlip, ToPILImage, Resize 14 | 15 | import core 16 | 17 | 18 | # ========== Set global settings ========== 19 | global_seed = 666 20 | deterministic = True 21 | torch.manual_seed(global_seed) 22 | # CUDA_SELECTED_DEVICES = '0,3' 23 | CUDA_SELECTED_DEVICES = '0' 24 | datasets_root_dir = '../datasets' 25 | 26 | 27 | # ========== BaselineMNISTNetwork_MNIST_BadNets ========== 28 | dataset = torchvision.datasets.MNIST 29 | 30 | transform_train = Compose([ 31 | ToTensor() 32 | ]) 33 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 34 | 35 | transform_test = Compose([ 36 | ToTensor() 37 | ]) 38 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 39 | 40 | pattern = torch.zeros((28, 28), dtype=torch.uint8) 41 | pattern[-3:, -3:] = 255 42 | weight = torch.zeros((28, 28), dtype=torch.float32) 43 | weight[-3:, -3:] = 1.0 44 | 45 | badnets = core.BadNets( 46 | train_dataset=trainset, 47 | test_dataset=testset, 48 | model=core.models.BaselineMNISTNetwork(), 49 | loss=nn.CrossEntropyLoss(), 50 | y_target=1, 51 | poisoned_rate=0.05, 52 | pattern=pattern, 53 | weight=weight, 54 | seed=global_seed, 55 | deterministic=deterministic 56 | ) 57 | 58 | # Train Attacked Model (schedule is set by yamengxi) 59 | schedule = { 60 | 'device': 'GPU', 61 | 'CUDA_SELECTED_DEVICES': CUDA_SELECTED_DEVICES, 62 | 63 | 'benign_training': False, 64 | 'batch_size': 1024, 65 | 'num_workers': 4, 66 | 67 | 'lr': 0.1, 68 | 'momentum': 0.9, 69 | 'weight_decay': 5e-4, 70 | 'gamma': 0.1, 71 | 'schedule': [150, 180], 72 | 73 | 'epochs': 200, 74 | 75 | 'log_iteration_interval': 100, 76 | 'test_epoch_interval': 10, 77 | 'save_epoch_interval': 10, 78 | 79 | 'save_dir': 'experiments', 80 | 'experiment_name': 'BaselineMNISTNetwork_MNIST_BadNets' 81 | } 82 | badnets.train(schedule) 83 | badnets.test(schedule) 84 | 85 | 86 | # ========== ResNet-18_CIFAR-10_BadNets ========== 87 | dataset = torchvision.datasets.CIFAR10 88 | 89 | transform_train = Compose([ 90 | RandomHorizontalFlip(), 91 | ToTensor() 92 | ]) 93 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 94 | 95 | transform_test = Compose([ 96 | ToTensor() 97 | ]) 98 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 99 | 100 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 101 | pattern[-3:, -3:] = 255 102 | weight = torch.zeros((32, 32), dtype=torch.float32) 103 | weight[-3:, -3:] = 1.0 104 | 105 | badnets = core.BadNets( 106 | train_dataset=trainset, 107 | test_dataset=testset, 108 | model=core.models.ResNet(18), 109 | loss=nn.CrossEntropyLoss(), 110 | y_target=1, 111 | poisoned_rate=0.05, 112 | pattern=pattern, 113 | weight=weight, 114 | seed=global_seed, 115 | deterministic=deterministic 116 | ) 117 | 118 | # Train Attacked Model (schedule is the same as https://github.com/THUYimingLi/Open-sourced_Dataset_Protection/blob/main/CIFAR/train_watermarked.py) 119 | schedule = { 120 | 'device': 'GPU', 121 | 'CUDA_SELECTED_DEVICES': CUDA_SELECTED_DEVICES, 122 | 123 | 'benign_training': False, 124 | 'batch_size': 1024, 125 | 'num_workers': 4, 126 | 127 | 'lr': 0.1, 128 | 'momentum': 0.9, 129 | 'weight_decay': 5e-4, 130 | 'gamma': 0.1, 131 | 'schedule': [150, 180], 132 | 133 | 'epochs': 200, 134 | 135 | 'log_iteration_interval': 100, 136 | 'test_epoch_interval': 10, 137 | 'save_epoch_interval': 10, 138 | 139 | 'save_dir': 'experiments', 140 | 'experiment_name': 'ResNet-18_CIFAR-10_BadNets' 141 | } 142 | badnets.train(schedule) 143 | badnets.test(schedule) 144 | 145 | 146 | # ========== ResNet-18_GTSRB_BadNets ========== 147 | transform_train = Compose([ 148 | ToPILImage(), 149 | Resize((32, 32)), 150 | ToTensor() 151 | ]) 152 | trainset = DatasetFolder( 153 | root=osp.join(datasets_root_dir, 'GTSRB', 'train'), # please replace this with path to your training set 154 | loader=cv2.imread, 155 | extensions=('png',), 156 | transform=transform_train, 157 | target_transform=None, 158 | is_valid_file=None) 159 | 160 | transform_test = Compose([ 161 | ToPILImage(), 162 | Resize((32, 32)), 163 | ToTensor() 164 | ]) 165 | testset = DatasetFolder( 166 | root=osp.join(datasets_root_dir, 'GTSRB', 'testset'), # please replace this with path to your test set 167 | loader=cv2.imread, 168 | extensions=('png',), 169 | transform=transform_test, 170 | target_transform=None, 171 | is_valid_file=None) 172 | 173 | 174 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 175 | pattern[-3:, -3:] = 255 176 | weight = torch.zeros((32, 32), dtype=torch.float32) 177 | weight[-3:, -3:] = 1.0 178 | 179 | badnets = core.BadNets( 180 | train_dataset=trainset, 181 | test_dataset=testset, 182 | model=core.models.ResNet(18, 43), 183 | loss=nn.CrossEntropyLoss(), 184 | y_target=1, 185 | poisoned_rate=0.05, 186 | pattern=pattern, 187 | weight=weight, 188 | poisoned_transform_train_index=2, 189 | poisoned_transform_test_index=2, 190 | seed=global_seed, 191 | deterministic=deterministic 192 | ) 193 | 194 | # Train Attacked Model (schedule is the same as https://github.com/THUYimingLi/Open-sourced_Dataset_Protection/blob/main/GTSRB/train_watermarked.py) 195 | schedule = { 196 | 'device': 'GPU', 197 | 'CUDA_SELECTED_DEVICES': CUDA_SELECTED_DEVICES, 198 | 199 | 'benign_training': False, 200 | 'batch_size': 1024, 201 | 'num_workers': 4, 202 | 203 | 'lr': 0.01, 204 | 'momentum': 0.9, 205 | 'weight_decay': 5e-4, 206 | 'gamma': 0.1, 207 | 'schedule': [20], 208 | 209 | 'epochs': 30, 210 | 211 | 'log_iteration_interval': 100, 212 | 'test_epoch_interval': 10, 213 | 'save_epoch_interval': 10, 214 | 215 | 'save_dir': 'experiments', 216 | 'experiment_name': 'ResNet-18_GTSRB_BadNets' 217 | } 218 | badnets.train(schedule) 219 | badnets.test(schedule) 220 | -------------------------------------------------------------------------------- /tests/test_BadNets_with_multi_gpus.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of poisoned training under BadNets. 3 | ''' 4 | 5 | 6 | import os.path as osp 7 | 8 | import cv2 9 | import torch 10 | import torch.nn as nn 11 | import torchvision 12 | from torchvision.datasets import DatasetFolder 13 | from torchvision.transforms import Compose, ToTensor, RandomHorizontalFlip, ToPILImage, Resize 14 | 15 | import core 16 | 17 | 18 | # ========== Set global settings ========== 19 | global_seed = 666 20 | deterministic = True 21 | torch.manual_seed(global_seed) 22 | CUDA_SELECTED_DEVICES = '0,3' 23 | # CUDA_SELECTED_DEVICES = '0' 24 | datasets_root_dir = '../datasets' 25 | 26 | 27 | # ========== BaselineMNISTNetwork_MNIST_BadNets ========== 28 | dataset = torchvision.datasets.MNIST 29 | 30 | transform_train = Compose([ 31 | ToTensor() 32 | ]) 33 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 34 | 35 | transform_test = Compose([ 36 | ToTensor() 37 | ]) 38 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 39 | 40 | pattern = torch.zeros((28, 28), dtype=torch.uint8) 41 | pattern[-3:, -3:] = 255 42 | weight = torch.zeros((28, 28), dtype=torch.float32) 43 | weight[-3:, -3:] = 1.0 44 | 45 | badnets = core.BadNets( 46 | train_dataset=trainset, 47 | test_dataset=testset, 48 | model=core.models.BaselineMNISTNetwork(), 49 | loss=nn.CrossEntropyLoss(), 50 | y_target=1, 51 | poisoned_rate=0.05, 52 | pattern=pattern, 53 | weight=weight, 54 | seed=global_seed, 55 | deterministic=deterministic 56 | ) 57 | 58 | # Train Attacked Model (schedule is set by yamengxi) 59 | schedule = { 60 | 'device': 'GPU', 61 | 'CUDA_SELECTED_DEVICES': CUDA_SELECTED_DEVICES, 62 | 63 | 'benign_training': False, 64 | 'batch_size': 1024, 65 | 'num_workers': 4, 66 | 67 | 'lr': 0.1, 68 | 'momentum': 0.9, 69 | 'weight_decay': 5e-4, 70 | 'gamma': 0.1, 71 | 'schedule': [150, 180], 72 | 73 | 'epochs': 200, 74 | 75 | 'log_iteration_interval': 100, 76 | 'test_epoch_interval': 10, 77 | 'save_epoch_interval': 10, 78 | 79 | 'save_dir': 'experiments', 80 | 'experiment_name': 'BaselineMNISTNetwork_MNIST_BadNets' 81 | } 82 | badnets.train(schedule) 83 | badnets.test(schedule) 84 | 85 | 86 | # ========== ResNet-18_CIFAR-10_BadNets ========== 87 | dataset = torchvision.datasets.CIFAR10 88 | 89 | transform_train = Compose([ 90 | RandomHorizontalFlip(), 91 | ToTensor() 92 | ]) 93 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 94 | 95 | transform_test = Compose([ 96 | ToTensor() 97 | ]) 98 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 99 | 100 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 101 | pattern[-3:, -3:] = 255 102 | weight = torch.zeros((32, 32), dtype=torch.float32) 103 | weight[-3:, -3:] = 1.0 104 | 105 | badnets = core.BadNets( 106 | train_dataset=trainset, 107 | test_dataset=testset, 108 | model=core.models.ResNet(18), 109 | loss=nn.CrossEntropyLoss(), 110 | y_target=1, 111 | poisoned_rate=0.05, 112 | pattern=pattern, 113 | weight=weight, 114 | seed=global_seed, 115 | deterministic=deterministic 116 | ) 117 | 118 | # Train Attacked Model (schedule is the same as https://github.com/THUYimingLi/Open-sourced_Dataset_Protection/blob/main/CIFAR/train_watermarked.py) 119 | schedule = { 120 | 'device': 'GPU', 121 | 'CUDA_SELECTED_DEVICES': CUDA_SELECTED_DEVICES, 122 | 123 | 'benign_training': False, 124 | 'batch_size': 1024, 125 | 'num_workers': 4, 126 | 127 | 'lr': 0.1, 128 | 'momentum': 0.9, 129 | 'weight_decay': 5e-4, 130 | 'gamma': 0.1, 131 | 'schedule': [150, 180], 132 | 133 | 'epochs': 200, 134 | 135 | 'log_iteration_interval': 100, 136 | 'test_epoch_interval': 10, 137 | 'save_epoch_interval': 10, 138 | 139 | 'save_dir': 'experiments', 140 | 'experiment_name': 'ResNet-18_CIFAR-10_BadNets' 141 | } 142 | badnets.train(schedule) 143 | badnets.test(schedule) 144 | 145 | 146 | # ========== ResNet-18_GTSRB_BadNets ========== 147 | transform_train = Compose([ 148 | ToPILImage(), 149 | Resize((32, 32)), 150 | ToTensor() 151 | ]) 152 | trainset = DatasetFolder( 153 | root=osp.join(datasets_root_dir, 'GTSRB', 'train'), # please replace this with path to your training set 154 | loader=cv2.imread, 155 | extensions=('png',), 156 | transform=transform_train, 157 | target_transform=None, 158 | is_valid_file=None) 159 | 160 | transform_test = Compose([ 161 | ToPILImage(), 162 | Resize((32, 32)), 163 | ToTensor() 164 | ]) 165 | testset = DatasetFolder( 166 | root=osp.join(datasets_root_dir, 'GTSRB', 'testset'), # please replace this with path to your test set 167 | loader=cv2.imread, 168 | extensions=('png',), 169 | transform=transform_test, 170 | target_transform=None, 171 | is_valid_file=None) 172 | 173 | 174 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 175 | pattern[-3:, -3:] = 255 176 | weight = torch.zeros((32, 32), dtype=torch.float32) 177 | weight[-3:, -3:] = 1.0 178 | 179 | badnets = core.BadNets( 180 | train_dataset=trainset, 181 | test_dataset=testset, 182 | model=core.models.ResNet(18, 43), 183 | loss=nn.CrossEntropyLoss(), 184 | y_target=1, 185 | poisoned_rate=0.05, 186 | pattern=pattern, 187 | weight=weight, 188 | poisoned_transform_train_index=2, 189 | poisoned_transform_test_index=2, 190 | seed=global_seed, 191 | deterministic=deterministic 192 | ) 193 | 194 | # Train Attacked Model (schedule is the same as https://github.com/THUYimingLi/Open-sourced_Dataset_Protection/blob/main/GTSRB/train_watermarked.py) 195 | schedule = { 196 | 'device': 'GPU', 197 | 'CUDA_SELECTED_DEVICES': CUDA_SELECTED_DEVICES, 198 | 199 | 'benign_training': False, 200 | 'batch_size': 1024, 201 | 'num_workers': 4, 202 | 203 | 'lr': 0.01, 204 | 'momentum': 0.9, 205 | 'weight_decay': 5e-4, 206 | 'gamma': 0.1, 207 | 'schedule': [20], 208 | 209 | 'epochs': 30, 210 | 211 | 'log_iteration_interval': 100, 212 | 'test_epoch_interval': 10, 213 | 'save_epoch_interval': 10, 214 | 215 | 'save_dir': 'experiments', 216 | 'experiment_name': 'ResNet-18_GTSRB_BadNets' 217 | } 218 | badnets.train(schedule) 219 | badnets.test(schedule) 220 | -------------------------------------------------------------------------------- /tests/test_Blended.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of benign training and poisoned training under Blended Attack. 3 | ''' 4 | 5 | 6 | import os 7 | 8 | import torch 9 | import torch.nn as nn 10 | from torch.utils.data import Dataset 11 | import torchvision 12 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 13 | import core 14 | 15 | 16 | global_seed = 666 17 | deterministic = True 18 | torch.manual_seed(global_seed) 19 | 20 | # Define Benign Training and Testing Dataset 21 | # dataset = torchvision.datasets.CIFAR10 22 | dataset = torchvision.datasets.MNIST 23 | 24 | 25 | transform_train = Compose([ 26 | ToTensor(), 27 | RandomHorizontalFlip() 28 | ]) 29 | trainset = dataset('data', train=True, transform=transform_train, download=True) 30 | 31 | transform_test = Compose([ 32 | ToTensor() 33 | ]) 34 | testset = dataset('data', train=False, transform=transform_test, download=True) 35 | 36 | 37 | # Show an Example of Benign Training Samples 38 | index = 44 39 | 40 | x, y = trainset[index] 41 | print(y) 42 | for a in x[0]: 43 | for b in a: 44 | print("%-4.2f" % float(b), end=' ') 45 | print() 46 | 47 | 48 | # Settings of Pattern and Weight 49 | ''' 50 | pattern = torch.zeros((1, 32, 32), dtype=torch.uint8) 51 | pattern[0, -3:, -3:] = 255 52 | weight = torch.zeros((1, 32, 32), dtype=torch.float32) 53 | weight[0, -3:, -3:] = 0.2 54 | ''' 55 | pattern = torch.zeros((1, 28, 28), dtype=torch.uint8) 56 | pattern[0, -3:, -3:] = 255 57 | weight = torch.zeros((1, 28, 28), dtype=torch.float32) 58 | weight[0, -3:, -3:] = 0.2 59 | 60 | 61 | blended = core.Blended( 62 | train_dataset=trainset, 63 | test_dataset=testset, 64 | # model=core.models.ResNet(18), 65 | model=core.models.BaselineMNISTNetwork(), 66 | loss=nn.CrossEntropyLoss(), 67 | pattern=pattern, 68 | weight=weight, 69 | y_target=1, 70 | poisoned_rate=0.05, 71 | seed=global_seed, 72 | deterministic=deterministic 73 | ) 74 | 75 | poisoned_train_dataset, poisoned_test_dataset = blended.get_poisoned_dataset() 76 | 77 | 78 | # Show an Example of Poisoned Training Samples 79 | x, y = poisoned_train_dataset[index] 80 | print(y) 81 | for a in x[0]: 82 | for b in a: 83 | print("%-4.2f" % float(b), end=' ') 84 | print() 85 | 86 | 87 | # Show an Example of Poisoned Testing Samples 88 | x, y = poisoned_test_dataset[index] 89 | print(y) 90 | for a in x[0]: 91 | for b in a: 92 | print("%-4.2f" % float(b), end=' ') 93 | print() 94 | 95 | 96 | # Train Benign Model 97 | schedule = { 98 | 'device': 'GPU', 99 | 'CUDA_VISIBLE_DEVICES': '1', 100 | 'GPU_num': 1, 101 | 102 | 'benign_training': True, 103 | 'batch_size': 128, 104 | 'num_workers': 4, 105 | 106 | 'lr': 0.1, 107 | 'momentum': 0.9, 108 | 'weight_decay': 5e-4, 109 | 'gamma': 0.1, 110 | 'schedule': [150, 180], 111 | 112 | 'epochs': 200, 113 | 114 | 'log_iteration_interval': 100, 115 | 'test_epoch_interval': 10, 116 | 'save_epoch_interval': 10, 117 | 118 | 'save_dir': 'experiments', 119 | # 'experiment_name': 'train_benign_CIFAR10_Blended' 120 | 'experiment_name': 'train_benign_MNIST_Blended' 121 | } 122 | 123 | blended.train(schedule) 124 | benign_model = blended.get_model() 125 | 126 | 127 | # Test Benign Model 128 | test_schedule = { 129 | 'device': 'GPU', 130 | 'CUDA_VISIBLE_DEVICES': '1', 131 | 'GPU_num': 1, 132 | 133 | 'batch_size': 128, 134 | 'num_workers': 4, 135 | 136 | 'save_dir': 'experiments', 137 | # 'experiment_name': 'test_benign_CIFAR10_Blended' 138 | 'experiment_name': 'test_benign_MNIST_Blended' 139 | } 140 | blended.test(test_schedule) 141 | 142 | blended.model = core.models.BaselineMNISTNetwork() 143 | # Train Infected Model 144 | schedule = { 145 | 'device': 'GPU', 146 | 'CUDA_VISIBLE_DEVICES': '1', 147 | 'GPU_num': 1, 148 | 149 | 'benign_training': False, 150 | 'batch_size': 128, 151 | 'num_workers': 4, 152 | 153 | 'lr': 0.1, 154 | 'momentum': 0.9, 155 | 'weight_decay': 5e-4, 156 | 'gamma': 0.1, 157 | 'schedule': [150, 180], 158 | 159 | 'epochs': 200, 160 | 161 | 'log_iteration_interval': 100, 162 | 'test_epoch_interval': 10, 163 | 'save_epoch_interval': 10, 164 | 165 | 'save_dir': 'experiments', 166 | # 'experiment_name': 'train_poisoned_CIFAR10_Blended' 167 | 'experiment_name': 'train_poisoned_MNIST_Blended' 168 | } 169 | 170 | blended.train(schedule) 171 | infected_model = blended.get_model() 172 | 173 | 174 | # Test Infected Model 175 | test_schedule = { 176 | 'device': 'GPU', 177 | 'CUDA_VISIBLE_DEVICES': '1', 178 | 'GPU_num': 1, 179 | 180 | 'batch_size': 128, 181 | 'num_workers': 4, 182 | 183 | 'save_dir': 'experiments', 184 | # 'experiment_name': 'test_poisoned_CIFAR10_Blended' 185 | 'experiment_name': 'test_poisoned_MNIST_Blended' 186 | } 187 | blended.test(test_schedule) 188 | -------------------------------------------------------------------------------- /tests/test_Blind.py: -------------------------------------------------------------------------------- 1 | 2 | ''' 3 | This is the test code of poisoned training on GTSRB, CIFAR10, MNIST, using dataset class of torchvision.datasets.DatasetFolder torchvision.datasets.CIFAR10 torchvision.datasets.MNIST. 4 | The attack method is Blind. 5 | ''' 6 | 7 | 8 | import os 9 | from typing import Pattern 10 | import cv2 11 | import torch 12 | import torch.nn as nn 13 | from torch.utils.data import Dataset, dataloader 14 | import torchvision 15 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 16 | from torchvision import transforms 17 | from torch.utils.data import DataLoader 18 | from torchvision.datasets import DatasetFolder, CIFAR10, MNIST 19 | import core 20 | 21 | 22 | global_seed = 666 23 | deterministic = True 24 | torch.manual_seed(global_seed) 25 | 26 | alpha = torch.zeros(3,32,32) 27 | alpha[:,-3:,-3:]=1. 28 | pattern = torch.zeros(3,32,32) 29 | pattern[:,-3:,-3:]=1. 30 | 31 | def show_dataset(dataset, num, path_to_save): 32 | """Each image in dataset should be torch.Tensor, shape (C,H,W)""" 33 | import matplotlib.pyplot as plt 34 | plt.figure(figsize=(20,20)) 35 | for i in range(num): 36 | ax = plt.subplot(num,1,i+1) 37 | img = (dataset[i][0]).permute(1,2,0).cpu().detach().numpy() 38 | ax.imshow(img) 39 | plt.savefig(path_to_save) 40 | 41 | # ===== Train backdoored model on GTSRB using with DatasetFolder ====== 42 | 43 | # Prepare datasets 44 | transform_train = Compose([ 45 | transforms.ToPILImage(), 46 | transforms.Resize((32, 32)), 47 | RandomHorizontalFlip(), 48 | ToTensor(), 49 | ]) 50 | transform_test = Compose([ 51 | transforms.ToPILImage(), 52 | transforms.Resize((32, 32)), 53 | ToTensor(), 54 | ]) 55 | 56 | trainset = DatasetFolder( 57 | root='/data/ganguanhao/datasets/GTSRB/train', # please replace this with path to your training set 58 | loader=cv2.imread, 59 | extensions=('png',), 60 | transform=transform_train, 61 | target_transform=None, 62 | is_valid_file=None) 63 | 64 | testset = DatasetFolder( 65 | root='/data/ganguanhao/datasets/GTSRB/testset', # please replace this with path to your test set 66 | loader=cv2.imread, 67 | extensions=('png',), 68 | transform=transform_train, 69 | target_transform=None, 70 | is_valid_file=None) 71 | 72 | # Configure the attack scheme 73 | blind= core.Blind( 74 | train_dataset=trainset, 75 | test_dataset=testset, 76 | model=core.models.ResNet(18, 43), 77 | loss=nn.CrossEntropyLoss(), 78 | y_target=1, 79 | pattern=pattern, 80 | alpha=alpha, 81 | schedule=None, 82 | seed=global_seed, 83 | deterministic=deterministic, 84 | use_neural_cleanse=False, 85 | ) 86 | 87 | 88 | schedule = { 89 | 'device': 'GPU', 90 | 'CUDA_VISIBLE_DEVICES': '0', 91 | 'GPU_num': 1, 92 | 93 | 'benign_training': False, 94 | 'batch_size': 128, 95 | 'num_workers': 8, 96 | 97 | 'lr': 0.1, 98 | 'momentum': 0.9, 99 | 'weight_decay': 5e-4, 100 | 'gamma': 0.1, 101 | 'schedule': [50, 75], 102 | 103 | 'epochs': 100, 104 | 105 | 'log_iteration_interval': 100, 106 | 'test_epoch_interval': 10, 107 | 'save_epoch_interval': 10, 108 | 109 | 'save_dir': 'experiments', 110 | 'experiment_name': 'train_poison_DataFolder_GTSRB_Blind' 111 | } 112 | 113 | # Train backdoored model 114 | blind.train(schedule) 115 | poisoned_trainset, poisoned_testset = blind.get_poisoned_dataset() 116 | show_dataset(poisoned_trainset, 5, 'gtsrb_train_poison.png') 117 | show_dataset(poisoned_testset, 5, 'gtsrb_test_poison.png') 118 | # # ===== Train backdoored model on GTSRB using with DatasetFolder (done) ====== 119 | 120 | # # ===== Train backdoored model on CIFAR10 using with CIFAR10 ===== 121 | 122 | # Prepare datasets 123 | transform_train = Compose([ 124 | transforms.Resize((32, 32)), 125 | RandomHorizontalFlip(), 126 | ToTensor(), 127 | # transforms.Normalize((0.485, 0.456, 0.406), 128 | # (0.229, 0.224, 0.225)) 129 | ]) 130 | transform_test = Compose([ 131 | transforms.Resize((32, 32)), 132 | ToTensor(), 133 | # transforms.Normalize((0.485, 0.456, 0.406), 134 | # (0.229, 0.224, 0.225)) 135 | ]) 136 | trainset = CIFAR10( 137 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 138 | transform=transform_train, 139 | target_transform=None, 140 | train=True, 141 | download=True) 142 | testset = CIFAR10( 143 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 144 | transform=transform_test, 145 | target_transform=None, 146 | train=False, 147 | download=True) 148 | 149 | 150 | blind= core.Blind( 151 | train_dataset=trainset, 152 | test_dataset=testset, 153 | model=core.models.ResNet(18), 154 | loss=nn.CrossEntropyLoss(), 155 | y_target=1, 156 | pattern=pattern, 157 | alpha=alpha, 158 | schedule=None, 159 | seed=global_seed, 160 | deterministic=deterministic, 161 | use_neural_cleanse=False, 162 | ) 163 | 164 | schedule = { 165 | 'device': 'GPU', 166 | 'CUDA_VISIBLE_DEVICES': '0', 167 | 'GPU_num': 1, 168 | 169 | 'benign_training': False, 170 | 'batch_size': 128, 171 | # 'batch_size': 64, 172 | 'num_workers': 8, 173 | 174 | 'lr': 0.1, 175 | 'momentum': 0.9, 176 | 'weight_decay': 5e-4, 177 | 'gamma': 0.1, 178 | 'schedule': [50, 75], 179 | 180 | 'epochs': 100, 181 | 182 | 'log_iteration_interval': 100, 183 | 'test_epoch_interval': 10, 184 | 'save_epoch_interval': 10, 185 | 186 | 'save_dir': 'experiments', 187 | 'experiment_name': 'train_poison_CIFAR10_Blind' 188 | } 189 | 190 | 191 | 192 | # Train backdoored model 193 | blind.train(schedule) 194 | poisoned_trainset, poisoned_testset = blind.get_poisoned_dataset() 195 | show_dataset(poisoned_trainset, 5, 'cifar_train_poison.png') 196 | show_dataset(poisoned_testset, 5, 'cifar_test_poison.png') 197 | 198 | 199 | 200 | 201 | # ===== Train backdoored model on CIFAR10 using with CIFAR10 (done)===== 202 | 203 | # ===== Train backdoored model on MNIST using with MNIST ===== 204 | # Prepare datasets 205 | transform_train = Compose([ 206 | transforms.Resize((28, 28)), 207 | RandomHorizontalFlip(), 208 | ToTensor(), 209 | ]) 210 | transform_test = Compose([ 211 | transforms.Resize((28, 28)), 212 | ToTensor(), 213 | ]) 214 | trainset = MNIST( 215 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 216 | transform=transform_train, 217 | target_transform=None, 218 | train=True, 219 | download=True) 220 | testset = MNIST( 221 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 222 | transform=transform_test, 223 | target_transform=None, 224 | train=False, 225 | download=True) 226 | 227 | loader = DataLoader(trainset,) 228 | 229 | alpha = torch.zeros(1,28,28) 230 | alpha[:,-3:,-3:]=1. 231 | pattern = torch.zeros(1,28,28) 232 | pattern[:,-3:,-3:]=1. 233 | # Configure the attack scheme 234 | blind= core.Blind( 235 | train_dataset=trainset, 236 | test_dataset=testset, 237 | model=core.models.BaselineMNISTNetwork(), 238 | loss=nn.CrossEntropyLoss(), 239 | y_target=1, 240 | pattern=pattern, 241 | alpha=alpha, 242 | schedule=None, 243 | seed=global_seed, 244 | deterministic=deterministic, 245 | use_neural_cleanse=False, 246 | ) 247 | 248 | schedule = { 249 | 'device': 'GPU', 250 | 'CUDA_VISIBLE_DEVICES': '0', 251 | 'GPU_num': 1, 252 | 253 | 'benign_training': False, 254 | 'batch_size': 128, 255 | 'num_workers': 8, 256 | 257 | 'lr': 0.1, 258 | 'momentum': 0.9, 259 | 'weight_decay': 5e-4, 260 | 'gamma': 0.1, 261 | 'schedule': [10, 15], 262 | 263 | 'epochs': 20, 264 | 265 | 'log_iteration_interval': 100, 266 | 'test_epoch_interval': 5, 267 | 'save_epoch_interval': 5, 268 | 269 | 'save_dir': 'experiments', 270 | 'experiment_name': 'train_poison_MNIST_Blind' 271 | } 272 | # Train backdoored model 273 | blind.train(schedule) 274 | poisoned_trainset, poisoned_testset = blind.get_poisoned_dataset() 275 | show_dataset(poisoned_trainset, 5, 'mnist_train_poison.png') 276 | show_dataset(poisoned_testset, 5, 'mnist_test_poison.png') 277 | # # ===== Train backdoored model on MNIST using with MNIST (done)===== 278 | -------------------------------------------------------------------------------- /tests/test_IBD-PSC.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of IBD-PSC defense. 3 | IBD-PSC: Input-level Backdoor Detection via Parameter-oriented Scaling Consistency [ICML, 2024] (https://arxiv.org/abs/2405.09786) 4 | 5 | ''' 6 | 7 | 8 | from copy import deepcopy 9 | import os.path as osp 10 | 11 | import numpy as np 12 | import random 13 | import cv2 14 | import os 15 | import torch 16 | import torch.nn as nn 17 | import torch.nn.functional as F 18 | 19 | import torchvision 20 | import torchvision.transforms as transforms 21 | from torchvision.datasets import DatasetFolder 22 | from torch.utils.data import Subset 23 | from torchvision.transforms import Compose, RandomHorizontalFlip, ToTensor, ToPILImage, Resize 24 | 25 | import core 26 | 27 | 28 | # ========== Set global settings ========== 29 | global_seed = 666 30 | # deterministic = True 31 | deterministic = False 32 | torch.manual_seed(global_seed) 33 | # ========== Set global settings ========== 34 | datasets_root_dir = os.path.expanduser('~/data/dataset') 35 | CUDA_VISIBLE_DEVICES = '0' 36 | portion = 0.1 37 | batch_size = 128 38 | num_workers = 4 39 | 40 | 41 | def test(model_name, dataset_name, attack_name, defense_name, benign_dataset, attacked_dataset, defense, y_target): 42 | schedule = { 43 | 'device': 'GPU', 44 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 45 | 'GPU_num': 1, 46 | 47 | 'batch_size': batch_size, 48 | 'num_workers': num_workers, 49 | 50 | 'metric': 'BA', 51 | 52 | 'save_dir': 'experiments/IBD-PSC-defense', 53 | 'experiment_name': f'{model_name}_{dataset_name}_{attack_name}_{defense_name}_BA' 54 | } 55 | defense.test_acc(benign_dataset, schedule) 56 | if not attack_name == 'Benign': 57 | schedule = { 58 | 'device': 'GPU', 59 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 60 | 'GPU_num': 1, 61 | 62 | 'batch_size': batch_size, 63 | 'num_workers': num_workers, 64 | 65 | # 1. ASR: the attack success rate calculated on all poisoned samples 66 | # 2. ASR_NoTarget: the attack success rate calculated on all poisoned samples whose ground-truth labels are not the target label 67 | # 3. BA: the accuracy on all benign samples 68 | # Hint: For ASR and BA, the computation of the metric is decided by the dataset but not schedule['metric']. 69 | # In other words, ASR or BA does not influence the computation of the metric. 70 | # For ASR_NoTarget, the code will delete all the samples whose ground-truth labels are the target label and then compute the metric. 71 | 'metric': 'ASR_NoTarget', 72 | 'y_target': y_target, 73 | 74 | 'save_dir': 'experiments/IBD-PSC-defense', 75 | 'experiment_name': f'{model_name}_{dataset_name}_{attack_name}_{defense_name}_ASR' 76 | } 77 | defense.test_acc(attacked_dataset, schedule) 78 | 79 | 80 | # ========== ResNet-18_CIFAR-10_Attack_IBD-PSC ========== 81 | # ===================== Benign ====================== 82 | model_name, dataset_name, attack_name, defense_name = 'ResNet-18', 'CIFAR-10', 'Benign', 'IBD-PSC' 83 | 84 | benign_model = core.models.ResNet(18, num_classes=10) 85 | model_path = '/home/ay2/data/experiments/List_ResNet-18_CIFAR-10_Benign/_2023-08-01_19:05:02/ckpt_999.pth' 86 | benign_model.load_state_dict(torch.load(model_path), strict=False) 87 | 88 | dataset = torchvision.datasets.CIFAR10 89 | transform_train = Compose([ 90 | RandomHorizontalFlip(), 91 | ToTensor() 92 | ]) 93 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 94 | transform_test = Compose([ 95 | ToTensor() 96 | ]) 97 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 98 | # Construct Shift Set for Defensive Purpose 99 | num_img = len(testset) 100 | indices = list(range(0, num_img)) 101 | random.shuffle(indices) 102 | val_budget = 2000 103 | val_indices = indices[:val_budget] 104 | # test_indices = indices[val_budget:] 105 | # # Construct Shift Set for Defensive Purpose 106 | # sub_testset = Subset(testset, test_indices) 107 | val_set = Subset(testset, val_indices) 108 | 109 | defense = core.IBD_PSC(model=benign_model, valset=val_set) 110 | test(model_name, dataset_name, attack_name, defense_name, testset, None, defense, None) 111 | 112 | # # ========== ResNet-18_CIFAR-10_Attack_Defense IBD-PSC ========== 113 | poisoning_rate = 0.1 114 | target_label = 0 115 | # ===================== BadNets ====================== 116 | model_name, dataset_name, attack_name, defense_name = 'ResNet-18', 'CIFAR-10', 'BadNets', 'IBD-PSC' 117 | badnet_model = core.models.ResNet(18, num_classes=10) 118 | model_path = '/home/ay2/data/experiments/CIFAR-10/ResNet-18/BadNets/pratio_0.1/2023-10-19_11:12:59/ckpt_0th_size_3_bottom-right_255_target_0.pth' 119 | def load_dict(model_path): 120 | state_dict = torch.load(model_path) 121 | # print(state_dict) 122 | if 'model' in list(state_dict.keys()): 123 | return state_dict['model'] 124 | else: 125 | return state_dict 126 | badnet_model.load_state_dict(load_dict(model_path)) 127 | 128 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 129 | pattern[-3:, -3:] = 255 130 | weight = torch.zeros((32, 32), dtype=torch.float32) 131 | weight[-3:, -3:] = 1.0 132 | 133 | # Get BadNets poisoned dataset 134 | attack = core.BadNets( 135 | train_dataset=trainset, 136 | test_dataset=testset, 137 | model=core.models.ResNet(18, num_classes=10), 138 | loss=nn.CrossEntropyLoss(), 139 | y_target=target_label, 140 | poisoned_rate=poisoning_rate, 141 | pattern=pattern, 142 | weight=weight, 143 | seed=global_seed, 144 | deterministic=deterministic 145 | ) 146 | poisoned_trainset, poisoned_testset = attack.get_poisoned_dataset() 147 | # poisoned_testset = Subset(poisoned_testset, test_indices) 148 | defense = core.IBD_PSC(model=badnet_model, valset=val_set) 149 | print(f'the BA and ASR of the original BadNets model: ............. ') 150 | test(model_name, dataset_name, attack_name, defense_name, testset, poisoned_testset, defense, None) 151 | defense.test(testset, poisoned_testset) 152 | 153 | # `detect` function is used to check if the first batch of samples in the dataset is poisoned. 154 | # Users can assemble their data into a batch of shape [num_samples, n, w, h] and call `defense._detect(batch)` 155 | # for online detection of the input. 156 | # `preds_benign` contains the detection results for the original test dataset. 157 | # `preds_poison` contains the detection results for the poisoned test dataset. 158 | preds_benign = defense.detect(testset) 159 | preds_poison = defense.detect(poisoned_testset) 160 | print(f'Is poisoned for real benign batch: {preds_benign}') 161 | print(f'Is poisoned for real poisoned batch: {preds_poison}') 162 | 163 | 164 | # ===================== WaNet ====================== 165 | attack_name = 'WaNet' 166 | wanet_model = core.models.ResNet(18, num_classes=10) 167 | model_path = '/home/ay2/data/experiments/CIFAR-10/ResNet-18/Wanet/ay3_wanet/0th_k_4_s_0.5_target_0_morph.pth.tar' 168 | state_dict = torch.load(model_path) 169 | wanet_model.load_state_dict(state_dict['netC']) 170 | identity_grid, noise_grid = state_dict["identity_grid"].cpu(), state_dict["noise_grid"].cpu() 171 | # target_label = state_dict['target'] 172 | 173 | attack = core.WaNet( 174 | train_dataset=trainset, 175 | test_dataset=testset, 176 | model=core.models.ResNet(18, num_classes=10), 177 | loss=nn.CrossEntropyLoss(), 178 | y_target=target_label, 179 | poisoned_rate=poisoning_rate, 180 | identity_grid=identity_grid, 181 | noise_grid=noise_grid, 182 | noise=True, 183 | seed=global_seed, 184 | deterministic=deterministic 185 | ) 186 | poisoned_trainset, poisoned_testset = attack.get_poisoned_dataset() 187 | defense = core.IBD_PSC(model=wanet_model, valset=val_set) 188 | print(f'the BA and ASR of the original WaNet model: ............. ') 189 | test(model_name, dataset_name, attack_name, defense_name, testset, poisoned_testset, defense, None) 190 | defense.test(testset, poisoned_testset) 191 | 192 | # `detect` function is used to check if the first batch of samples in the dataset is poisoned. 193 | # Users can assemble their data into a batch of shape [num_samples, n, w, h] and call `defense._detect(batch)` 194 | # for online detection of the input. 195 | # `preds_benign` contains the detection results for the original test dataset. 196 | # `preds_poison` contains the detection results for the poisoned test dataset. 197 | preds_benign = defense.detect(testset) 198 | preds_poison = defense.detect(poisoned_testset) 199 | print(f'Is poisoned for real benign batch: {preds_benign}') 200 | print(f'Is poisoned for real poisoned batch: {preds_poison}') 201 | 202 | -------------------------------------------------------------------------------- /tests/test_PhysicalBA.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of poisoned training under PhysicalBA. 3 | Using dataset class of torchvision.datasets.DatasetFolder, torchvision.datasets.MNIST and torchvision.datasets.CIFAR10. 4 | Default physical transformations is Compose([RandomHorizontalFlip(),ColorJitter(), RandomAffine()]) 5 | Choose other transformations from torchvsion.transforms if you need 6 | ''' 7 | 8 | 9 | import os 10 | import cv2 11 | import torch 12 | import torch.nn as nn 13 | from torch.utils.data import Dataset 14 | import torchvision 15 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip, ColorJitter, RandomAffine 16 | import torchvision.transforms as transforms 17 | import core 18 | 19 | global_seed = 666 20 | deterministic = True 21 | torch.manual_seed(global_seed) 22 | 23 | # ============== GTSRB ============== 24 | 25 | dataset = torchvision.datasets.DatasetFolder 26 | 27 | transform_train = Compose([ 28 | transforms.ToPILImage(), 29 | transforms.Resize((32, 32)), 30 | ToTensor() 31 | ]) 32 | trainset = dataset( 33 | root='/data/xutong/test/BackdoorBox/tests/data/GTSRB/train', 34 | loader=cv2.imread, 35 | extensions=('png',), 36 | transform=transform_train, 37 | target_transform=None, 38 | is_valid_file=None) 39 | 40 | transform_test = Compose([ 41 | transforms.ToPILImage(), 42 | transforms.Resize((32, 32)), 43 | ToTensor() 44 | ]) 45 | testset = dataset( 46 | root='/data/xutong/test/BackdoorBox/tests/data/GTSRB/testset', 47 | loader=cv2.imread, 48 | extensions=('png',), 49 | transform=transform_test, 50 | target_transform=None, 51 | is_valid_file=None) 52 | 53 | pattern = torch.zeros((1, 32, 32), dtype=torch.uint8) 54 | pattern[0, -3:, -3:] = 255 55 | weight = torch.zeros((1, 32, 32), dtype=torch.float32) 56 | weight[0, -3:, -3:] = 1.0 57 | 58 | PhysicalBA = core.PhysicalBA( 59 | train_dataset=trainset, 60 | test_dataset=testset, 61 | model=core.models.ResNet(18,43), 62 | loss=nn.CrossEntropyLoss(), 63 | y_target=1, 64 | poisoned_rate=0.05, 65 | pattern=pattern, 66 | weight=weight, 67 | poisoned_transform_train_index=2, 68 | poisoned_transform_test_index=2, 69 | schedule=None, 70 | seed=666, 71 | # modify other transformations from torchvsion.transforms if you want 72 | physical_transformations = Compose([ 73 | RandomHorizontalFlip(), 74 | ColorJitter(brightness=0.2,contrast=0.2), 75 | RandomAffine(degrees=10,translate=(0.1, 0.1), scale=(0.8, 0.9)) 76 | ]) 77 | ) 78 | 79 | poisoned_train_dataset, poisoned_test_dataset = PhysicalBA.get_poisoned_dataset() 80 | 81 | # train attacked model 82 | schedule = { 83 | 'device': 'GPU', 84 | 'CUDA_VISIBLE_DEVICES': '0', 85 | 'GPU_num': 1, 86 | 87 | 'benign_training': False, 88 | 'batch_size': 128, 89 | 'num_workers': 16, 90 | 91 | 'lr': 0.1, 92 | 'momentum': 0.9, 93 | 'weight_decay': 5e-4, 94 | 'gamma': 0.1, 95 | 'schedule': [150, 180], 96 | 97 | 'epochs': 200, 98 | 99 | 'log_iteration_interval': 100, 100 | 'test_epoch_interval': 10, 101 | 'save_epoch_interval': 10, 102 | 103 | 'save_dir': 'experiments', 104 | 'experiment_name': 'train_poisoned_DatasetFolder_PhysicalBA' 105 | } 106 | 107 | PhysicalBA.train(schedule) 108 | infected_model = PhysicalBA.get_model() 109 | 110 | # Test Infected Model 111 | test_schedule = { 112 | 'device': 'GPU', 113 | 'CUDA_VISIBLE_DEVICES': '0', 114 | 'GPU_num': 1, 115 | 116 | 'batch_size': 128, 117 | 'num_workers': 16, 118 | 119 | 'save_dir': 'experiments', 120 | 'experiment_name': 'test_poisoned_DatasetFolder_PhysicalBA' 121 | } 122 | PhysicalBA.test(test_schedule) 123 | 124 | # ============== mnist ============== 125 | 126 | dataset = torchvision.datasets.MNIST 127 | 128 | transform_train = Compose([ 129 | ToTensor(), 130 | ]) 131 | trainset = dataset('data', train=True, transform=transform_train, download=True) 132 | 133 | transform_test = Compose([ 134 | ToTensor() 135 | ]) 136 | testset = dataset('data', train=False, transform=transform_test, download=True) 137 | 138 | PhysicalBA = core.PhysicalBA( 139 | train_dataset=trainset, 140 | test_dataset=testset, 141 | model=core.models.BaselineMNISTNetwork(), 142 | loss=nn.CrossEntropyLoss(), 143 | y_target=1, 144 | poisoned_rate=0.05, 145 | seed=global_seed, 146 | deterministic=deterministic, 147 | # modify other transformations from torchvsion.transforms if you want 148 | physical_transformations = Compose([ 149 | RandomHorizontalFlip(), 150 | ColorJitter(brightness=0.2), 151 | RandomAffine(degrees=10,translate=(0.1, 0.1), scale=(0.8, 0.9)) 152 | ]) 153 | ) 154 | 155 | poisoned_train_dataset, poisoned_test_dataset = PhysicalBA.get_poisoned_dataset() 156 | 157 | # Train Infected Model 158 | schedule = { 159 | 'device': 'GPU', 160 | 'CUDA_VISIBLE_DEVICES': '0', 161 | 'GPU_num': 1, 162 | 163 | 'benign_training': False, # Train Infected Model 164 | 'batch_size': 128, 165 | 'num_workers': 4, 166 | 167 | 'lr': 0.1, 168 | 'momentum': 0.9, 169 | 'weight_decay': 5e-4, 170 | 'gamma': 0.1, 171 | 'schedule': [150, 180], 172 | 173 | 'epochs': 200, 174 | 175 | 'log_iteration_interval': 100, 176 | 'test_epoch_interval': 10, 177 | 'save_epoch_interval': 10, 178 | 179 | 'save_dir': 'experiments', 180 | 'experiment_name': 'train_poisoned_MNIST_PhysicalBA' 181 | } 182 | 183 | PhysicalBA.train(schedule) 184 | infected_model = PhysicalBA.get_model() 185 | 186 | # Test Infected Model 187 | test_schedule = { 188 | 'device': 'GPU', 189 | 'CUDA_VISIBLE_DEVICES': '0', 190 | 'GPU_num': 1, 191 | 192 | 'batch_size': 128, 193 | 'num_workers': 4, 194 | 195 | 'save_dir': 'experiments', 196 | 'experiment_name': 'test_poisoned_MNIST_PhysicalBA' 197 | } 198 | PhysicalBA.test(test_schedule) 199 | 200 | # ============== cifar10 ============== 201 | 202 | dataset = torchvision.datasets.CIFAR10 203 | 204 | transform_train = Compose([ 205 | ToTensor(), 206 | ]) 207 | trainset = dataset('data', train=True, transform=transform_train, download=True) 208 | 209 | transform_test = Compose([ 210 | ToTensor() 211 | ]) 212 | testset = dataset('data', train=False, transform=transform_test, download=True) 213 | 214 | PhysicalBA = core.PhysicalBA( 215 | train_dataset=trainset, 216 | test_dataset=testset, 217 | model=core.models.ResNet(18), 218 | loss=nn.CrossEntropyLoss(), 219 | y_target=1, 220 | poisoned_rate=0.05, 221 | seed=global_seed, 222 | deterministic=deterministic, 223 | # modify other transformations from torchvsion.transforms if you want 224 | physical_transformations = Compose([ 225 | RandomHorizontalFlip(), 226 | ColorJitter(brightness=0.2,contrast=0.2), 227 | RandomAffine(degrees=10,translate=(0.1, 0.1), scale=(0.8, 0.9)) 228 | ]) 229 | ) 230 | 231 | poisoned_train_dataset, poisoned_test_dataset = PhysicalBA.get_poisoned_dataset() 232 | 233 | # Train Infected Model 234 | schedule = { 235 | 'device': 'GPU', 236 | 'CUDA_VISIBLE_DEVICES': '0', 237 | 'GPU_num': 1, 238 | 239 | 'benign_training': False, # Train Infected Model 240 | 'batch_size': 128, 241 | 'num_workers': 4, 242 | 243 | 'lr': 0.1, 244 | 'momentum': 0.9, 245 | 'weight_decay': 5e-4, 246 | 'gamma': 0.1, 247 | 'schedule': [150, 180], 248 | 249 | 'epochs': 200, 250 | 251 | 'log_iteration_interval': 100, 252 | 'test_epoch_interval': 10, 253 | 'save_epoch_interval': 10, 254 | 255 | 'save_dir': 'experiments', 256 | 'experiment_name': 'train_poisoned_CIFAR10_PhysicalBA' 257 | } 258 | 259 | PhysicalBA.train(schedule) 260 | infected_model = PhysicalBA.get_model() 261 | 262 | # Test Infected Model 263 | test_schedule = { 264 | 'device': 'GPU', 265 | 'CUDA_VISIBLE_DEVICES': '0', 266 | 'GPU_num': 1, 267 | 268 | 'batch_size': 128, 269 | 'num_workers': 4, 270 | 271 | 'save_dir': 'experiments', 272 | 'experiment_name': 'test_poisoned_CIFAR10_PhysicalBA' 273 | } 274 | PhysicalBA.test(test_schedule) -------------------------------------------------------------------------------- /tests/test_REFINE.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of REFINE defense. You need to train an input transformation module before testing. 3 | You can refer to `train_REFINE.py` for more details. 4 | ''' 5 | 6 | import sys 7 | import os 8 | sys.path.append(os.getcwd()) 9 | import torch 10 | import core 11 | import argparse 12 | from getmodel import GetModel 13 | from getdataset import GetDataset 14 | 15 | parser = argparse.ArgumentParser(description='PyTorch ShrinkPad') 16 | parser.add_argument('--gpu_id', default='0', type=str, help='id(s) for CUDA_VISIBLE_DEVICES') 17 | parser.add_argument('--dataset', default='CIFAR10', type=str) 18 | parser.add_argument('--model', default='ResNet18', type=str) 19 | parser.add_argument('--attack', default='BadNets', type=str) 20 | parser.add_argument('--tlabel', default=0, type=int) 21 | 22 | args = parser.parse_args() 23 | state = {k: v for k, v in args._get_kwargs()} 24 | 25 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id 26 | 27 | # ========== Set training model ========== 28 | getmodel = GetModel(args.dataset, args.model, args.attack) 29 | my_model = getmodel.get_model() 30 | 31 | # ========== Set global settings ========== 32 | global_seed = 666 33 | deterministic = True 34 | torch.manual_seed(global_seed) 35 | CUDA_VISIBLE_DEVICES = args.gpu_id 36 | save_path = f'REFINE/{args.dataset}/{args.model}/{args.attack}' 37 | if not os.path.exists(save_path): 38 | os.makedirs(save_path) 39 | 40 | batch_size = 128 41 | num_workers = 4 42 | 43 | getdataset = GetDataset(args.dataset, args.model, args.attack, args.tlabel) 44 | trainset, testset = getdataset.get_benign_dataset() 45 | poisoned_trainset, poisoned_testset = getdataset.get_poisoned_dataset() 46 | 47 | if args.dataset == 'CIFAR10': 48 | size_map = 32 49 | first_channel = 64 50 | elif args.dataset == 'ImageNet50': 51 | size_map = 224 52 | first_channel = 32 53 | 54 | # If you have trained 'unet' and defined 'arr_shuffle' in core.REFINE, you can also get them (instead of loading them directly locally) as follows: 55 | # unet_trained = defense.unet 56 | # arr_shuffle_defined = defense.arr_shuffle 57 | 58 | defense = core.REFINE( 59 | unet=core.models.UNetLittle(args=None, n_channels=3, n_classes=3, first_channels=first_channel), 60 | pretrain='/data/home/Yukun/BackdoorBox/REFINE/CIFAR10/ResNet18/BadNets/REFINE_train_2025-02-14_17:24:44/ckpt_epoch_150.pth', 61 | arr_path='/data/home/Yukun/BackdoorBox/REFINE/CIFAR10/ResNet18/BadNets/REFINE_train_2025-02-14_17:24:44/label_shuffle.pth', 62 | model=my_model, 63 | num_classes=10, 64 | seed=global_seed, 65 | deterministic=deterministic 66 | ) 67 | 68 | schedule = { 69 | 'device': 'GPU', 70 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 71 | 'GPU_num': 1, 72 | 73 | # 'test_model': model_path, 74 | 'batch_size': batch_size, 75 | 'num_workers': num_workers, 76 | 77 | 'metric': 'BA', 78 | 79 | 'save_dir': save_path, 80 | 'experiment_name': f'REFINE_BA' 81 | } 82 | defense.test(testset, schedule) 83 | 84 | schedule = { 85 | 'device': 'GPU', 86 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 87 | 'GPU_num': 1, 88 | 89 | # 'test_model': model_path, 90 | 'batch_size': batch_size, 91 | 'num_workers': num_workers, 92 | 93 | # 1. ASR: the attack success rate calculated on all poisoned samples 94 | # 2. ASR_NoTarget: the attack success rate calculated on all poisoned samples whose ground-truth labels are not the target label 95 | # 3. BA: the accuracy on all benign samples 96 | # Hint: For ASR and BA, the computation of the metric is decided by the dataset but not schedule['metric']. 97 | # In other words, ASR or BA does not influence the computation of the metric. 98 | # For ASR_NoTarget, the code will delete all the samples whose ground-truth labels are the target label and then compute the metric. 99 | 'metric': 'ASR_NoTarget', 100 | 'y_target': args.tlabel, 101 | 102 | 'save_dir': save_path, 103 | 'experiment_name': f'REFINE_ASR' 104 | } 105 | defense.test(poisoned_testset, schedule) 106 | -------------------------------------------------------------------------------- /tests/test_SCALE_UP.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of SCALE-UP defense. 3 | SCALE-UP: An Efficient Black-box Input-level Backdoor Detection via Analyzing Scaled Prediction Consistency [ICLR, 2023] (https://arxiv.org/abs/2302.03251) 4 | ''' 5 | 6 | 7 | from copy import deepcopy 8 | import os.path as osp 9 | 10 | import numpy as np 11 | import random 12 | import cv2 13 | import os 14 | import torch 15 | import torch.nn as nn 16 | import torch.nn.functional as F 17 | 18 | import torchvision 19 | import torchvision.transforms as transforms 20 | from torchvision.datasets import DatasetFolder 21 | from torch.utils.data import Subset 22 | from torchvision.transforms import Compose, RandomHorizontalFlip, ToTensor, ToPILImage, Resize 23 | from PIL import Image 24 | import core 25 | 26 | 27 | # ========== Set global settings ========== 28 | global_seed = 0 29 | # deterministic = True 30 | deterministic = False 31 | torch.manual_seed(global_seed) 32 | # ========== Set global settings ========== 33 | datasets_root_dir = os.path.expanduser('~/data/dataset') 34 | CUDA_VISIBLE_DEVICES = '0' 35 | portion = 0.1 36 | batch_size = 128 37 | num_workers = 4 38 | 39 | 40 | def test(model_name, dataset_name, attack_name, defense_name, benign_dataset, attacked_dataset, defense, y_target): 41 | schedule = { 42 | 'device': 'GPU', 43 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 44 | 'GPU_num': 1, 45 | 46 | 'batch_size': batch_size, 47 | 'num_workers': num_workers, 48 | 49 | 'metric': 'BA', 50 | 51 | 'save_dir': 'experiments/SCALE-UP-defense', 52 | 'experiment_name': f'{model_name}_{dataset_name}_{attack_name}_{defense_name}_BA' 53 | } 54 | defense.test_acc(benign_dataset, schedule) 55 | if not attack_name == 'Benign': 56 | schedule = { 57 | 'device': 'GPU', 58 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 59 | 'GPU_num': 1, 60 | 61 | 'batch_size': batch_size, 62 | 'num_workers': num_workers, 63 | 64 | # 1. ASR: the attack success rate calculated on all poisoned samples 65 | # 2. ASR_NoTarget: the attack success rate calculated on all poisoned samples whose ground-truth labels are not the target label 66 | # 3. BA: the accuracy on all benign samples 67 | # Hint: For ASR and BA, the computation of the metric is decided by the dataset but not schedule['metric']. 68 | # In other words, ASR or BA does not influence the computation of the metric. 69 | # For ASR_NoTarget, the code will delete all the samples whose ground-truth labels are the target label and then compute the metric. 70 | 'metric': 'ASR_NoTarget', 71 | 'y_target': y_target, 72 | 73 | 'save_dir': 'experiments/SCALE-UP-defense', 74 | 'experiment_name': f'{model_name}_{dataset_name}_{attack_name}_{defense_name}_ASR' 75 | } 76 | defense.test_acc(attacked_dataset, schedule) 77 | 78 | 79 | # ========== ResNet-18_CIFAR-10_Attack_SCALE-UP ========== 80 | # ===================== Benign ====================== 81 | model_name, dataset_name, attack_name, defense_name = 'ResNet-18', 'CIFAR-10', 'Benign', 'SCALE-UP' 82 | 83 | benign_model = core.models.ResNet(18, num_classes=10) 84 | model_path = '/home/ay2/data/experiments/List_ResNet-18_CIFAR-10_Benign/_2023-08-01_19:05:02/ckpt_999.pth' 85 | benign_model.load_state_dict(torch.load(model_path), strict=False) 86 | 87 | dataset = torchvision.datasets.CIFAR10 88 | transform_train = Compose([ 89 | RandomHorizontalFlip(), 90 | ToTensor() 91 | ]) 92 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 93 | transform_test = Compose([ 94 | ToTensor() 95 | ]) 96 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 97 | # Construct Shift Set for Defensive Purpose 98 | num_img = len(testset) 99 | indices = list(range(0, num_img)) 100 | random.shuffle(indices) 101 | val_budget = 2000 102 | val_indices = indices[:val_budget] 103 | test_indices = indices[val_budget:] 104 | # Construct Shift Set for Defensive Purpose 105 | sub_testset = Subset(testset, test_indices) 106 | val_set = Subset(testset, val_indices) 107 | 108 | defense = core.SCALE_UP(model=benign_model) 109 | test(model_name, dataset_name, attack_name, defense_name, testset, None, defense, None) 110 | 111 | # # ========== ResNet-18_CIFAR-10_Attack_Defense SCALE-UP ========== 112 | poisoning_rate = 0.1 113 | target_label = 0 114 | # ===================== BadNets ====================== 115 | model_name, dataset_name, attack_name, defense_name = 'ResNet-18', 'CIFAR-10', 'BadNets', 'SCALE-UP' 116 | badnet_model = core.models.ResNet(18, num_classes=10) 117 | model_path = '/home/ay2/data/experiments/CIFAR10/ResNet-18/BadNets/pratio_0.1/2024-05-23_16:07:13/ckpt_ random.pth' 118 | def load_dict(model_path): 119 | state_dict = torch.load(model_path) 120 | # print(state_dict) 121 | if 'model' in list(state_dict.keys()): 122 | return state_dict['model'] 123 | else: 124 | return state_dict 125 | badnet_model.load_state_dict(load_dict(model_path)) 126 | 127 | pattern_path = '/home/ay2/data/experiments/CIFAR10/ResNet-18/BadNets/pratio_0.1/2024-05-23_16:07:13/triggers/random_pattern.pt' 128 | pattern = torch.load(pattern_path) 129 | weight = torch.zeros((32, 32), dtype=torch.float32) 130 | weight[-3:, -3:] = 1.0 131 | # Get BadNets poisoned dataset 132 | attack = core.BadNets( 133 | train_dataset=trainset, 134 | test_dataset=testset, 135 | model=core.models.ResNet(18, num_classes=10), 136 | loss=nn.CrossEntropyLoss(), 137 | y_target=target_label, 138 | poisoned_rate=poisoning_rate, 139 | pattern=pattern, 140 | weight=weight, 141 | seed=global_seed, 142 | deterministic=deterministic 143 | ) 144 | poisoned_trainset, poisoned_testset = attack.get_poisoned_dataset() 145 | # poisoned_testset = Subset(poisoned_testset, test_indices) 146 | defense = core.SCALE_UP(model=badnet_model) 147 | print(f'the BA and ASR of the original BadNets model: ............. ') 148 | # test(model_name, dataset_name, attack_name, defense_name, testset, poisoned_testset, defense, None) 149 | print('---------data-free scenario----------') 150 | defense.test(testset, poisoned_testset) 151 | print('---------data-limited scenario----------') 152 | defense_with_val = core.SCALE_UP(model=badnet_model, valset=val_set) 153 | defense_with_val.test(testset, poisoned_testset) 154 | 155 | # `detect` function is used to check if the first batch of samples in the dataset is poisoned. 156 | # Users can assemble their data into a batch of shape [num_samples, n, w, h] and call `defense._detect(batch)` 157 | # for online detection of the input. 158 | # `preds_benign` contains the detection results for the original test dataset. 159 | # `preds_poison` contains the detection results for the poisoned test dataset. 160 | preds_benign = defense.detect(testset) 161 | preds_poison = defense.detect(poisoned_testset) 162 | print(f'Is poisoned for real benign batch: {preds_benign}') 163 | print(f'Is poisoned for real poisoned batch: {preds_poison}') 164 | 165 | 166 | # ===================== WaNet ====================== 167 | attack_name = 'WaNet' 168 | wanet_model = core.models.ResNet(18, num_classes=10) 169 | model_path = '/home/ay2/data/experiments/CIFAR-10/ResNet-18/Wanet/ay3_wanet/0th_k_4_s_0.5_target_0_morph.pth.tar' 170 | state_dict = torch.load(model_path) 171 | wanet_model.load_state_dict(state_dict['netC']) 172 | identity_grid, noise_grid = state_dict["identity_grid"].cpu(), state_dict["noise_grid"].cpu() 173 | # target_label = state_dict['target'] 174 | 175 | attack = core.WaNet( 176 | train_dataset=trainset, 177 | test_dataset=testset, 178 | model=core.models.ResNet(18, num_classes=10), 179 | loss=nn.CrossEntropyLoss(), 180 | y_target=target_label, 181 | poisoned_rate=poisoning_rate, 182 | identity_grid=identity_grid, 183 | noise_grid=noise_grid, 184 | noise=True, 185 | seed=global_seed, 186 | deterministic=deterministic 187 | ) 188 | poisoned_testset = attack.poisoned_test_dataset 189 | defense = core.SCALE_UP(model=wanet_model) 190 | print(f'the BA and ASR of the original WaNet model: ............. ') 191 | test(model_name, dataset_name, attack_name, defense_name, testset, poisoned_testset, defense, None) 192 | print('---------data-free scenario----------') 193 | defense.test(testset, poisoned_testset) 194 | print('---------data-limited scenario----------') 195 | defense_with_val = core.SCALE_UP(model=wanet_model, valset=val_set) 196 | defense_with_val.test(testset, poisoned_testset) 197 | 198 | # `detect` function is used to check if the first batch of samples in the dataset is poisoned. 199 | # Users can assemble their data into a batch of shape [num_samples, n, w, h] and call `defense._detect(batch)` 200 | # for online detection of the input. 201 | # `preds_benign` contains the detection results for the original test dataset. 202 | # `preds_poison` contains the detection results for the poisoned test dataset. 203 | preds_benign = defense.detect(testset) 204 | preds_poison = defense.detect(poisoned_testset) 205 | print(f'Is poisoned for real benign batch: {preds_benign}') 206 | print(f'Is poisoned for real poisoned batch: {preds_poison}') 207 | -------------------------------------------------------------------------------- /tests/test_TUAP.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of poisoned training under TUAP. 3 | ''' 4 | 5 | 6 | import os 7 | import os.path as osp 8 | 9 | import cv2 10 | import numpy as np 11 | import torch 12 | import torch.nn as nn 13 | from torch.utils.data import Dataset 14 | import torchvision 15 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 16 | import torchvision.transforms as transforms 17 | from torchvision.datasets import DatasetFolder 18 | 19 | import core 20 | from core.attacks.TUAP import TUAP 21 | 22 | 23 | CUDA_VISIBLE_DEVICES = '3' 24 | os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_VISIBLE_DEVICES 25 | datasets_root_dir = '/data/yamengxi/Backdoor/datasets/' 26 | global_seed = 666 27 | deterministic = True 28 | torch.manual_seed(global_seed) 29 | 30 | 31 | # ============== TUAP attack BaselineMNISTNetwork on MNIST ============== 32 | dataset = torchvision.datasets.MNIST 33 | 34 | transform_train = Compose([ 35 | ToTensor() 36 | ]) 37 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 38 | print("trainset",len(trainset)) 39 | transform_test = Compose([ 40 | ToTensor() 41 | ]) 42 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 43 | print("testset",len(testset)) 44 | 45 | 46 | schedule = { 47 | 'device': 'GPU', 48 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 49 | 'GPU_num': 1, 50 | 51 | 'benign_training': False, # Train Attacked Model 52 | 'batch_size': 1024, 53 | 'num_workers': 8, 54 | 55 | 'lr': 0.1, 56 | 'momentum': 0.9, 57 | 'weight_decay': 5e-4, 58 | 'gamma': 0.1, 59 | 'schedule': [150, 180], 60 | 61 | 'epochs': 200, 62 | 63 | 'log_iteration_interval': 100, 64 | 'test_epoch_interval': 10, 65 | 'save_epoch_interval': 10, 66 | 67 | 'save_dir': 'experiments', 68 | 'experiment_name': 'BaselineMNISTNetwork_MNIST_TUAP' 69 | } 70 | 71 | 72 | UAP_benign_model = core.models.BaselineMNISTNetwork() 73 | UAP_benign_PATH = './benign_MNIST/MNIST/ckpt_epoch_20.pth' 74 | checkpoint = torch.load(UAP_benign_PATH) 75 | UAP_benign_model.load_state_dict(checkpoint) 76 | 77 | poisoned_rate = 0.25 78 | # epsilon = 76.0/255 79 | epsilon = 0.3 80 | delta = 0.2 81 | max_iter_uni = np.inf 82 | p_norm = np.inf 83 | num_classes = 10 84 | overshoot = 0.02 85 | max_iter_df = 50 86 | p_samples = 0.01 87 | mask = np.ones((1, 28, 28)) 88 | 89 | tuap = core.TUAP( 90 | train_dataset=trainset, 91 | test_dataset=testset, 92 | model=core.models.BaselineMNISTNetwork(), 93 | loss=nn.CrossEntropyLoss(), 94 | 95 | benign_model=UAP_benign_model, 96 | y_target=2, 97 | poisoned_rate=poisoned_rate, 98 | epsilon = epsilon, 99 | delta=delta, 100 | max_iter_uni=max_iter_uni, 101 | p_norm=p_norm, 102 | num_classes=num_classes, 103 | overshoot=overshoot, 104 | max_iter_df=max_iter_df, 105 | p_samples=p_samples, 106 | mask=mask, 107 | 108 | poisoned_transform_train_index=0, 109 | poisoned_transform_test_index=0, 110 | poisoned_target_transform_index=0, 111 | schedule=schedule, 112 | seed=global_seed, 113 | deterministic=True 114 | ) 115 | tuap.train() 116 | 117 | 118 | # ==============TUAP attack ResNet-18 on CIFAR-10 ============== 119 | dataset = torchvision.datasets.CIFAR10 120 | transform_train = Compose([ 121 | ToTensor(), 122 | RandomHorizontalFlip() 123 | ]) 124 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 125 | print("trainset",len(trainset)) 126 | transform_test = Compose([ 127 | ToTensor() 128 | ]) 129 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 130 | print("testset",len(testset)) 131 | 132 | 133 | schedule = { 134 | 'device': 'GPU', 135 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 136 | 'GPU_num': 1, 137 | 138 | 'benign_training': False, # Train Attacked Model 139 | 'batch_size': 128, 140 | 'num_workers': 8, 141 | 142 | 'lr': 0.1, 143 | 'momentum': 0.9, 144 | 'weight_decay': 5e-4, 145 | 'gamma': 0.1, 146 | 'schedule': [150, 180], 147 | 148 | 'epochs': 200, 149 | 150 | 'log_iteration_interval': 100, 151 | 'test_epoch_interval': 10, 152 | 'save_epoch_interval': 10, 153 | 154 | 'save_dir': 'experiments', 155 | 'experiment_name': 'ResNet-18_CIFAR-10_TUAP' 156 | } 157 | 158 | UAP_benign_model = core.models.ResNet(18) 159 | UAP_benign_PATH = './benign_CIFAR10/ckpt_epoch_100.pth' 160 | checkpoint = torch.load(UAP_benign_PATH) 161 | UAP_benign_model.load_state_dict(checkpoint) 162 | poisoned_rate = 0.25 163 | # epsilon = 10 164 | epsilon = 0.031 165 | delta = 0.2 166 | max_iter_uni = np.inf 167 | p_norm = np.inf 168 | num_classes = 10 169 | overshoot = 0.02 170 | max_iter_df = 50 171 | p_samples = 0.01 172 | mask = np.ones((3, 32, 32)) 173 | 174 | 175 | tuap = core.TUAP( 176 | train_dataset=trainset, 177 | test_dataset=testset, 178 | model=core.models.ResNet(18), 179 | loss=nn.CrossEntropyLoss(), 180 | 181 | benign_model=UAP_benign_model, 182 | y_target=2, 183 | poisoned_rate=poisoned_rate, 184 | epsilon = epsilon, 185 | delta=delta, 186 | max_iter_uni=max_iter_uni, 187 | p_norm=p_norm, 188 | num_classes=num_classes, 189 | overshoot=overshoot, 190 | max_iter_df=max_iter_df, 191 | p_samples=p_samples, 192 | mask=mask, 193 | 194 | poisoned_transform_train_index=0, 195 | poisoned_transform_test_index=0, 196 | poisoned_target_transform_index=0, 197 | schedule=schedule, 198 | seed=global_seed, 199 | deterministic=True 200 | ) 201 | 202 | tuap.train() 203 | 204 | 205 | 206 | # ============== TUAP attack ResNet-18 on GTSRB ============== 207 | transform_train = Compose([ 208 | transforms.ToPILImage(), 209 | transforms.Resize((32, 32)), 210 | RandomHorizontalFlip(), 211 | ToTensor() 212 | ]) 213 | trainset = DatasetFolder( 214 | root=osp.join(datasets_root_dir, 'GTSRB', 'train'), # please replace this with path to your training set 215 | loader=cv2.imread, 216 | extensions=('png',), 217 | transform=transform_train, 218 | target_transform=None, 219 | is_valid_file=None) 220 | 221 | transform_test = Compose([ 222 | transforms.ToPILImage(), 223 | transforms.Resize((32, 32)), 224 | ToTensor() 225 | ]) 226 | testset = DatasetFolder( 227 | root=osp.join(datasets_root_dir, 'GTSRB', 'testset'), # please replace this with path to your test set 228 | loader=cv2.imread, 229 | extensions=('png',), 230 | transform=transform_test, 231 | target_transform=None, 232 | is_valid_file=None) 233 | 234 | 235 | schedule = { 236 | 'device': 'GPU', 237 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 238 | 'GPU_num': 1, 239 | 240 | 'benign_training': False, # Train Attacked Model 241 | 'batch_size': 256, 242 | 'num_workers': 8, 243 | 244 | 'lr': 0.01, 245 | 'momentum': 0.9, 246 | 'weight_decay': 5e-4, 247 | 'gamma': 0.1, 248 | 'schedule': [20], 249 | 250 | 'epochs': 50, 251 | 252 | 'log_iteration_interval': 100, 253 | 'test_epoch_interval': 10, 254 | 'save_epoch_interval': 10, 255 | 256 | 'save_dir': 'experiments', 257 | 'experiment_name': 'ResNet-18_GTSRB_TUAP' 258 | } 259 | 260 | UAP_benign_model = core.models.ResNet(18,43) 261 | UAP_benign_PATH = './benign_GTSRB/ckpt_epoch_100.pth' 262 | checkpoint = torch.load(UAP_benign_PATH) 263 | UAP_benign_model.load_state_dict(checkpoint) 264 | poisoned_rate = 0.5 265 | epsilon = 0.031 266 | delta = 0.2 267 | max_iter_uni =20 268 | p_norm = np.inf 269 | num_classes = 10 270 | overshoot = 0.02 271 | max_iter_df = 50 272 | p_samples = 0.02 273 | mask = np.ones((3, 32, 32)) 274 | 275 | 276 | tuap = core.TUAP( 277 | train_dataset=trainset, 278 | test_dataset=testset, 279 | model=core.models.ResNet(18, 43), 280 | loss=nn.CrossEntropyLoss(), 281 | 282 | benign_model=UAP_benign_model, 283 | y_target=2, 284 | poisoned_rate=poisoned_rate, 285 | epsilon = epsilon, 286 | delta=delta, 287 | max_iter_uni=max_iter_uni, 288 | p_norm=p_norm, 289 | num_classes=num_classes, 290 | overshoot=overshoot, 291 | max_iter_df=max_iter_df, 292 | p_samples=p_samples, 293 | mask=mask, 294 | 295 | poisoned_transform_train_index=2, 296 | poisoned_transform_test_index=2, 297 | poisoned_target_transform_index=0, 298 | schedule=schedule, 299 | seed=global_seed, 300 | deterministic=True 301 | ) 302 | 303 | tuap.train() 304 | -------------------------------------------------------------------------------- /tests/test_cifar10.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of benign training and poisoned training on torchvision.datasets.CIFAR10. 3 | Attack method is BadNets. 4 | ''' 5 | 6 | 7 | import os 8 | 9 | import torch 10 | import torch.nn as nn 11 | from torch.utils.data import Dataset 12 | import torchvision 13 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 14 | 15 | import core 16 | 17 | # dataset = torchvision.datasets.MNIST 18 | dataset = torchvision.datasets.CIFAR10 19 | 20 | 21 | transform_train = Compose([ 22 | ToTensor(), 23 | RandomHorizontalFlip() 24 | ]) 25 | trainset = dataset('data', train=True, transform=transform_train, download=True) 26 | 27 | transform_test = Compose([ 28 | ToTensor() 29 | ]) 30 | testset = dataset('data', train=False, transform=transform_test, download=True) 31 | 32 | 33 | index = 44 34 | 35 | x, y = trainset[index] 36 | print(y) 37 | for a in x[0]: 38 | for b in a: 39 | print("%-4.2f" % float(b), end=' ') 40 | print() 41 | 42 | 43 | badnets = core.BadNets( 44 | train_dataset=trainset, 45 | test_dataset=testset, 46 | model=core.models.ResNet(18), 47 | # model=core.models.BaselineMNISTNetwork(), 48 | loss=nn.CrossEntropyLoss(), 49 | y_target=1, 50 | poisoned_rate=0.05, 51 | seed=666 52 | ) 53 | 54 | poisoned_train_dataset, poisoned_test_dataset = badnets.get_poisoned_dataset() 55 | 56 | 57 | x, y = poisoned_train_dataset[index] 58 | print(y) 59 | for a in x[0]: 60 | for b in a: 61 | print("%-4.2f" % float(b), end=' ') 62 | print() 63 | 64 | 65 | x, y = poisoned_test_dataset[index] 66 | print(y) 67 | for a in x[0]: 68 | for b in a: 69 | print("%-4.2f" % float(b), end=' ') 70 | print() 71 | 72 | # train benign model 73 | schedule = { 74 | 'device': 'GPU', 75 | 'CUDA_VISIBLE_DEVICES': '2', 76 | 'GPU_num': 1, 77 | 78 | 'benign_training': True, 79 | 'batch_size': 128, 80 | 'num_workers': 16, 81 | 82 | 'lr': 0.1, 83 | 'momentum': 0.9, 84 | 'weight_decay': 5e-4, 85 | 'gamma': 0.1, 86 | 'schedule': [150, 180], 87 | 88 | 'epochs': 200, 89 | 90 | 'log_iteration_interval': 100, 91 | 'test_epoch_interval': 10, 92 | 'save_epoch_interval': 10, 93 | 94 | 'save_dir': 'experiments', 95 | 'experiment_name': 'train_benign_CIFAR10' 96 | } 97 | 98 | badnets.train(schedule) 99 | 100 | # train attacked model 101 | schedule = { 102 | 'device': 'GPU', 103 | 'CUDA_VISIBLE_DEVICES': '2', 104 | 'GPU_num': 1, 105 | 106 | 'benign_training': False, 107 | 'batch_size': 128, 108 | 'num_workers': 16, 109 | 110 | 'lr': 0.1, 111 | 'momentum': 0.9, 112 | 'weight_decay': 5e-4, 113 | 'gamma': 0.1, 114 | 'schedule': [150, 180], 115 | 116 | 'epochs': 200, 117 | 118 | 'log_iteration_interval': 100, 119 | 'test_epoch_interval': 10, 120 | 'save_epoch_interval': 10, 121 | 122 | 'save_dir': 'experiments', 123 | 'experiment_name': 'train_poisoned_CIFAR10' 124 | } 125 | 126 | badnets.train(schedule) 127 | -------------------------------------------------------------------------------- /tests/test_dataset_folder.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of benign training and poisoned training on torchvision.datasets.DatasetFolder. 3 | Dataset is CIFAR-10. 4 | Attack method is BadNets. 5 | ''' 6 | 7 | 8 | import os 9 | 10 | import cv2 11 | import torch 12 | import torch.nn as nn 13 | from torch.utils.data import Dataset 14 | import torchvision 15 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 16 | 17 | import core 18 | 19 | 20 | dataset = torchvision.datasets.DatasetFolder 21 | 22 | # image file -> cv.imread -> numpy.ndarray (H x W x C) -> ToTensor -> torch.Tensor (C x H x W) -> RandomHorizontalFlip -> torch.Tensor -> network input 23 | transform_train = Compose([ 24 | ToTensor(), 25 | RandomHorizontalFlip() 26 | ]) 27 | trainset = dataset( 28 | root='/home/public/yamengxi/data/cifar10/train', 29 | loader=cv2.imread, 30 | extensions=('png',), 31 | transform=transform_train, 32 | target_transform=None, 33 | is_valid_file=None) 34 | 35 | transform_test = Compose([ 36 | ToTensor() 37 | ]) 38 | testset = dataset( 39 | root='/home/public/yamengxi/data/cifar10/test', 40 | loader=cv2.imread, 41 | extensions=('png',), 42 | transform=transform_train, 43 | target_transform=None, 44 | is_valid_file=None) 45 | 46 | index = 44 47 | 48 | x, y = trainset[index] 49 | print(y) 50 | for a in x[0]: 51 | for b in a: 52 | print("%-4.2f" % float(b), end=' ') 53 | print() 54 | 55 | 56 | pattern = torch.zeros((1, 32, 32), dtype=torch.uint8) 57 | pattern[0, -3:, -3:] = 255 58 | weight = torch.zeros((1, 32, 32), dtype=torch.float32) 59 | weight[0, -3:, -3:] = 1.0 60 | 61 | badnets = core.BadNets( 62 | train_dataset=trainset, 63 | test_dataset=testset, 64 | model=core.models.ResNet(18), 65 | # model=core.models.BaselineMNISTNetwork(), 66 | loss=nn.CrossEntropyLoss(), 67 | y_target=1, 68 | poisoned_rate=0.05, 69 | pattern=pattern, 70 | weight=weight, 71 | poisoned_transform_index=0, 72 | poisoned_target_transform_index=0, 73 | schedule=None, 74 | seed=666 75 | ) 76 | 77 | poisoned_train_dataset, poisoned_test_dataset = badnets.get_poisoned_dataset() 78 | 79 | x, y = poisoned_train_dataset[index] 80 | print(y) 81 | for a in x[0]: 82 | for b in a: 83 | print("%-4.2f" % float(b), end=' ') 84 | print() 85 | 86 | x, y = poisoned_test_dataset[index] 87 | print(y) 88 | for a in x[0]: 89 | for b in a: 90 | print("%-4.2f" % float(b), end=' ') 91 | print() 92 | 93 | # train benign model 94 | schedule = { 95 | 'device': 'GPU', 96 | 'CUDA_VISIBLE_DEVICES': '1', 97 | 'GPU_num': 1, 98 | 99 | 'benign_training': True, 100 | 'batch_size': 128, 101 | 'num_workers': 16, 102 | 103 | 'lr': 0.1, 104 | 'momentum': 0.9, 105 | 'weight_decay': 5e-4, 106 | 'gamma': 0.1, 107 | 'schedule': [150, 180], 108 | 109 | 'epochs': 200, 110 | 111 | 'log_iteration_interval': 100, 112 | 'test_epoch_interval': 10, 113 | 'save_epoch_interval': 10, 114 | 115 | 'save_dir': 'experiments', 116 | 'experiment_name': 'train_benign_DatasetFolder-CIFAR10' 117 | } 118 | 119 | badnets.train(schedule) 120 | 121 | # train attacked model 122 | schedule = { 123 | 'device': 'GPU', 124 | 'CUDA_VISIBLE_DEVICES': '1', 125 | 'GPU_num': 1, 126 | 127 | 'benign_training': False, 128 | 'batch_size': 128, 129 | 'num_workers': 16, 130 | 131 | 'lr': 0.1, 132 | 'momentum': 0.9, 133 | 'weight_decay': 5e-4, 134 | 'gamma': 0.1, 135 | 'schedule': [150, 180], 136 | 137 | 'epochs': 200, 138 | 139 | 'log_iteration_interval': 100, 140 | 'test_epoch_interval': 10, 141 | 'save_epoch_interval': 10, 142 | 143 | 'save_dir': 'experiments', 144 | 'experiment_name': 'train_poisoned_DatasetFolder-CIFAR10' 145 | } 146 | 147 | badnets.train(schedule) -------------------------------------------------------------------------------- /tests/test_refool.py: -------------------------------------------------------------------------------- 1 | 2 | ''' 3 | This is the test code of poisoned training on GTSRB, CIFAR10, MNIST, using dataset class of torchvision.datasets.DatasetFolder torchvision.datasets.CIFAR10 torchvision.datasets.MNIST. 4 | The attack method is Refool. 5 | ''' 6 | 7 | 8 | import os 9 | import cv2 10 | import torch 11 | import torch.nn as nn 12 | from torch.utils.data import Dataset, dataloader 13 | import torchvision 14 | from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip 15 | from torchvision import transforms 16 | from torch.utils.data import DataLoader 17 | from torchvision.datasets import DatasetFolder, CIFAR10, MNIST 18 | import core 19 | 20 | 21 | global_seed = 666 22 | deterministic = True 23 | torch.manual_seed(global_seed) 24 | 25 | # load reflection images 26 | reflection_images = [] 27 | reflection_data_dir = "/data/ganguanhao/datasets/VOCdevkit/VOC2012/JPEGImages/" # please replace this with path to your desired reflection set 28 | 29 | def read_image(img_path, type=None): 30 | img = cv2.imread(img_path) 31 | if type is None: 32 | return img 33 | elif isinstance(type,str) and type.upper() == "RGB": 34 | return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 35 | elif isinstance(type,str) and type.upper() == "GRAY": 36 | return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) 37 | else: 38 | raise NotImplementedError 39 | 40 | reflection_image_path = os.listdir(reflection_data_dir) 41 | reflection_images = [read_image(os.path.join(reflection_data_dir,img_path)) for img_path in reflection_image_path[:200]] 42 | 43 | 44 | # ===== Train backdoored model on GTSRB using with DatasetFolder ====== 45 | 46 | # Prepare datasets 47 | transform_train = Compose([ 48 | transforms.ToPILImage(), 49 | transforms.Resize((32, 32)), 50 | RandomHorizontalFlip(), 51 | ToTensor(), 52 | transforms.Normalize((0.485, 0.456, 0.406), 53 | (0.229, 0.224, 0.225)) 54 | ]) 55 | transform_test = Compose([ 56 | transforms.ToPILImage(), 57 | transforms.Resize((32, 32)), 58 | ToTensor(), 59 | transforms.Normalize((0.485, 0.456, 0.406), 60 | (0.229, 0.224, 0.225)) 61 | ]) 62 | 63 | trainset = DatasetFolder( 64 | root='/data/ganguanhao/datasets/GTSRB/train', # please replace this with path to your training set 65 | loader=cv2.imread, 66 | extensions=('png',), 67 | transform=transform_train, 68 | target_transform=None, 69 | is_valid_file=None) 70 | 71 | testset = DatasetFolder( 72 | root='/data/ganguanhao/datasets/GTSRB/testset', # please replace this with path to your test set 73 | loader=cv2.imread, 74 | extensions=('png',), 75 | transform=transform_test, 76 | target_transform=None, 77 | is_valid_file=None) 78 | 79 | # Configure the attack scheme 80 | refool= core.Refool( 81 | train_dataset=trainset, 82 | test_dataset=testset, 83 | model=core.models.ResNet(18, 43), 84 | loss=nn.CrossEntropyLoss(), 85 | y_target=1, 86 | poisoned_rate=0.05, 87 | poisoned_transform_train_index=0, 88 | poisoned_transform_test_index=0, 89 | poisoned_target_transform_index=0, 90 | schedule=None, 91 | seed=global_seed, 92 | deterministic=deterministic, 93 | reflection_candidates = reflection_images, 94 | ) 95 | 96 | 97 | schedule = { 98 | 'device': 'GPU', 99 | 'CUDA_VISIBLE_DEVICES': '0', 100 | 'GPU_num': 1, 101 | 102 | 'benign_training': False, 103 | 'batch_size': 128, 104 | 'num_workers': 8, 105 | 106 | 'lr': 0.1, 107 | 'momentum': 0.9, 108 | 'weight_decay': 5e-4, 109 | 'gamma': 0.1, 110 | 'schedule': [50, 75], 111 | 112 | 'epochs': 100, 113 | 114 | 'log_iteration_interval': 100, 115 | 'test_epoch_interval': 10, 116 | 'save_epoch_interval': 10, 117 | 118 | 'save_dir': 'experiments', 119 | 'experiment_name': 'train_poison_DataFolder_GTSRB_Refool' 120 | } 121 | 122 | # Train backdoored model 123 | refool.train(schedule) 124 | 125 | # ===== Train backdoored model on GTSRB using with DatasetFolder (done) ====== 126 | 127 | # ===== Train backdoored model on CIFAR10 using with CIFAR10 ===== 128 | 129 | # Prepare datasets 130 | transform_train = Compose([ 131 | transforms.Resize((32, 32)), 132 | RandomHorizontalFlip(), 133 | ToTensor(), 134 | transforms.Normalize((0.485, 0.456, 0.406), 135 | (0.229, 0.224, 0.225)) 136 | ]) 137 | transform_test = Compose([ 138 | transforms.Resize((32, 32)), 139 | ToTensor(), 140 | transforms.Normalize((0.485, 0.456, 0.406), 141 | (0.229, 0.224, 0.225)) 142 | ]) 143 | trainset = CIFAR10( 144 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 145 | transform=transform_train, 146 | target_transform=None, 147 | train=True, 148 | download=True) 149 | testset = CIFAR10( 150 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 151 | transform=transform_test, 152 | target_transform=None, 153 | train=False, 154 | download=True) 155 | 156 | 157 | # Configure the attack scheme 158 | refool= core.Refool( 159 | train_dataset=trainset, 160 | test_dataset=testset, 161 | model=core.models.ResNet(18), 162 | loss=nn.CrossEntropyLoss(), 163 | y_target=1, 164 | poisoned_rate=0.05, 165 | poisoned_transform_train_index=0, 166 | poisoned_transform_test_index=0, 167 | poisoned_target_transform_index=0, 168 | schedule=None, 169 | seed=global_seed, 170 | deterministic=deterministic, 171 | reflection_candidates = reflection_images, 172 | ) 173 | 174 | 175 | schedule = { 176 | 'device': 'GPU', 177 | 'CUDA_VISIBLE_DEVICES': '0', 178 | 'GPU_num': 1, 179 | 180 | 'benign_training': False, 181 | 'batch_size': 128, 182 | 'num_workers': 8, 183 | 184 | 'lr': 0.1, 185 | 'momentum': 0.9, 186 | 'weight_decay': 5e-4, 187 | 'gamma': 0.1, 188 | 'schedule': [50, 75], 189 | 190 | 'epochs': 100, 191 | 192 | 'log_iteration_interval': 100, 193 | 'test_epoch_interval': 10, 194 | 'save_epoch_interval': 10, 195 | 196 | 'save_dir': 'experiments', 197 | 'experiment_name': 'train_poison_CIFAR10_Refool' 198 | } 199 | 200 | 201 | # Train backdoored model 202 | refool.train(schedule) 203 | 204 | # ===== Train backdoored model on CIFAR10 using with CIFAR10 (done)===== 205 | 206 | # ===== Train backdoored model on MNIST using with MNIST ===== 207 | # Prepare datasets 208 | transform_train = Compose([ 209 | transforms.Resize((28, 28)), 210 | RandomHorizontalFlip(), 211 | ToTensor(), 212 | ]) 213 | transform_test = Compose([ 214 | transforms.Resize((28, 28)), 215 | ToTensor(), 216 | ]) 217 | trainset = MNIST( 218 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 219 | transform=transform_train, 220 | target_transform=None, 221 | train=True, 222 | download=True) 223 | testset = MNIST( 224 | root='/data/ganguanhao/datasets', # please replace this with path to your dataset 225 | transform=transform_test, 226 | target_transform=None, 227 | train=False, 228 | download=True) 229 | 230 | loader = DataLoader(trainset,) 231 | 232 | 233 | # Configure the attack scheme 234 | refool= core.Refool( 235 | train_dataset=trainset, 236 | test_dataset=testset, 237 | model=core.models.BaselineMNISTNetwork(), 238 | loss=nn.CrossEntropyLoss(), 239 | y_target=1, 240 | poisoned_rate=0.05, 241 | poisoned_transform_train_index=0, 242 | poisoned_transform_test_index=0, 243 | poisoned_target_transform_index=0, 244 | schedule=None, 245 | seed=global_seed, 246 | deterministic=deterministic, 247 | reflection_candidates = reflection_images, 248 | ) 249 | 250 | schedule = { 251 | 'device': 'GPU', 252 | 'CUDA_VISIBLE_DEVICES': '0', 253 | 'GPU_num': 1, 254 | 255 | 'benign_training': False, 256 | 'batch_size': 128, 257 | 'num_workers': 8, 258 | 259 | 'lr': 0.1, 260 | 'momentum': 0.9, 261 | 'weight_decay': 5e-4, 262 | 'gamma': 0.1, 263 | 'schedule': [10, 15], 264 | 265 | 'epochs': 20, 266 | 267 | 'log_iteration_interval': 100, 268 | 'test_epoch_interval': 5, 269 | 'save_epoch_interval': 5, 270 | 271 | 'save_dir': 'experiments', 272 | 'experiment_name': 'train_poison_MNIST_Refool' 273 | } 274 | # Train backdoored model 275 | refool.train(schedule) 276 | # ===== Train backdoored model on MNIST using with MNIST (done)===== 277 | -------------------------------------------------------------------------------- /tests/test_refool_vis.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "'''\n", 10 | "This is the test code for visualizing poisoned training on CIFAR10, MNIST, using dataset class of torchvision.datasets.DatasetFolder torchvision.datasets.CIFAR10 torchvision.datasets.MNIST.\n", 11 | "Attack method is Refool.\n", 12 | "'''\n", 13 | "\n", 14 | "import os\n", 15 | "import cv2\n", 16 | "import torch\n", 17 | "import torch.nn as nn\n", 18 | "from torch.utils.data import Dataset, dataloader\n", 19 | "import torchvision\n", 20 | "from torchvision import transforms\n", 21 | "from torchvision.datasets import CIFAR10, MNIST, DatasetFolder\n", 22 | "from torchvision.transforms import Compose, ToTensor, PILToTensor, RandomHorizontalFlip\n", 23 | "from torch.utils.data import DataLoader\n", 24 | "import core\n", 25 | "import matplotlib.pyplot as plt" 26 | ] 27 | }, 28 | { 29 | "cell_type": "code", 30 | "execution_count": null, 31 | "metadata": {}, 32 | "outputs": [], 33 | "source": [ 34 | "global_seed = 666\n", 35 | "deterministic = True\n", 36 | "torch.manual_seed(global_seed)\n", 37 | "reflection_data_dir = \"/data/ganguanhao/datasets/VOCdevkit/VOC2012/JPEGImages/\"\n", 38 | "def read_image(img_path, type=None):\n", 39 | " img = cv2.imread(img_path)\n", 40 | " if type is None: \n", 41 | " return img\n", 42 | " elif isinstance(type,str) and type.upper() == \"RGB\":\n", 43 | " return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n", 44 | " elif isinstance(type,str) and type.upper() == \"GRAY\":\n", 45 | " return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n", 46 | " else:\n", 47 | " raise NotImplementedError\n", 48 | " \n", 49 | "reflection_image_path = os.listdir(reflection_data_dir)\n", 50 | "reflection_images = [read_image(os.path.join(reflection_data_dir,img_path)) for img_path in reflection_image_path[:200]]\n", 51 | "\n", 52 | "def show_img(x):\n", 53 | " img = cv2.cvtColor(x.permute(1,2,0).numpy(),cv2.COLOR_BGR2RGB)\n", 54 | " plt.imshow(img)\n", 55 | " plt.show()" 56 | ] 57 | }, 58 | { 59 | "cell_type": "code", 60 | "execution_count": null, 61 | "metadata": {}, 62 | "outputs": [], 63 | "source": [ 64 | "transform_train = Compose([\n", 65 | " transforms.ToPILImage(),\n", 66 | " # transforms.Resize((32,32)),\n", 67 | " RandomHorizontalFlip(),\n", 68 | " ToTensor(),\n", 69 | "])\n", 70 | "transform_test = Compose([\n", 71 | " transforms.ToPILImage(),\n", 72 | " # transforms.Resize((32,32)),\n", 73 | " ToTensor(),\n", 74 | "])\n", 75 | "\n", 76 | "trainset = DatasetFolder(\n", 77 | " root='/data/ganguanhao/datasets/cifar10-folder/trainset',\n", 78 | " loader=cv2.imread,\n", 79 | " extensions=('png',),\n", 80 | " transform=transform_train,\n", 81 | " target_transform=None,\n", 82 | " is_valid_file=None)\n", 83 | "\n", 84 | "testset = DatasetFolder(\n", 85 | " root='/data/ganguanhao/datasets/cifar10-folder/testset',\n", 86 | " loader=cv2.imread,\n", 87 | " extensions=('png',),\n", 88 | " transform=transform_train,\n", 89 | " target_transform=None,\n", 90 | " is_valid_file=None)\n", 91 | "\n", 92 | "refool= core.Refool(\n", 93 | " train_dataset=trainset,\n", 94 | " test_dataset=testset,\n", 95 | " model=core.models.ResNet(18),\n", 96 | " loss=nn.CrossEntropyLoss(),\n", 97 | " y_target=1,\n", 98 | " poisoned_rate=0.05,\n", 99 | " poisoned_transform_train_index=0,\n", 100 | " poisoned_transform_test_index=0,\n", 101 | " poisoned_target_transform_index=0,\n", 102 | " schedule=None,\n", 103 | " seed=global_seed,\n", 104 | " deterministic=deterministic,\n", 105 | " reflection_candidates = reflection_images,\n", 106 | ")\n", 107 | "\n", 108 | "\n", 109 | "poisoned_train_dataset, poisoned_test_dataset = refool.get_poisoned_dataset()\n", 110 | "print(\"============train_images============\")\n", 111 | "for i in range(10):\n", 112 | " show_img(poisoned_train_dataset[i][0])\n", 113 | "print(\"============test_images============\")\n", 114 | "for i in range(10):\n", 115 | " show_img(poisoned_test_dataset[i][0])\n" 116 | ] 117 | }, 118 | { 119 | "cell_type": "code", 120 | "execution_count": null, 121 | "metadata": {}, 122 | "outputs": [], 123 | "source": [ 124 | "transform_train = Compose([\n", 125 | " RandomHorizontalFlip(),\n", 126 | " ToTensor(),\n", 127 | "])\n", 128 | "transform_test = Compose([\n", 129 | " ToTensor(),\n", 130 | "])\n", 131 | "trainset = CIFAR10(\n", 132 | " root='/data/ganguanhao/datasets',\n", 133 | " transform=transform_train,\n", 134 | " target_transform=None,\n", 135 | " train=True,\n", 136 | " download=True)\n", 137 | "testset = CIFAR10(\n", 138 | " root='/data/ganguanhao/datasets',\n", 139 | " transform=transform_test,\n", 140 | " target_transform=None,\n", 141 | " train=False,\n", 142 | " download=True)\n", 143 | "\n", 144 | "refool= core.Refool(\n", 145 | " train_dataset=trainset,\n", 146 | " test_dataset=testset,\n", 147 | " model=core.models.ResNet(18),\n", 148 | " loss=nn.CrossEntropyLoss(),\n", 149 | " y_target=1,\n", 150 | " poisoned_rate=0.05,\n", 151 | " poisoned_transform_train_index=0,\n", 152 | " poisoned_transform_test_index=0,\n", 153 | " poisoned_target_transform_index=0,\n", 154 | " schedule=None,\n", 155 | " seed=global_seed,\n", 156 | " deterministic=deterministic,\n", 157 | " reflection_candidates = reflection_images,\n", 158 | ")\n", 159 | "\n", 160 | "\n", 161 | "poisoned_train_dataset, poisoned_test_dataset = refool.get_poisoned_dataset()\n", 162 | "print(\"============train_images============\")\n", 163 | "for i in range(10):\n", 164 | " show_img(poisoned_train_dataset[i][0])\n", 165 | "print(\"============test_images============\")\n", 166 | "for i in range(10):\n", 167 | " show_img(poisoned_test_dataset[i][0])" 168 | ] 169 | }, 170 | { 171 | "cell_type": "code", 172 | "execution_count": null, 173 | "metadata": {}, 174 | "outputs": [], 175 | "source": [ 176 | "transform_train = Compose([\n", 177 | " RandomHorizontalFlip(),\n", 178 | " ToTensor(),\n", 179 | "])\n", 180 | "transform_test = Compose([\n", 181 | " ToTensor(),\n", 182 | "])\n", 183 | "trainset = MNIST(\n", 184 | " root='/data/ganguanhao/datasets',\n", 185 | " transform=transform_train,\n", 186 | " target_transform=None,\n", 187 | " train=True,\n", 188 | " download=True)\n", 189 | "testset = MNIST(\n", 190 | " root='/data/ganguanhao/datasets',\n", 191 | " transform=transform_test,\n", 192 | " target_transform=None,\n", 193 | " train=False,\n", 194 | " download=True)\n", 195 | "\n", 196 | "loader = DataLoader(trainset,)\n", 197 | "\n", 198 | "refool= core.Refool(\n", 199 | " train_dataset=trainset,\n", 200 | " test_dataset=testset,\n", 201 | " model=core.models.BaselineMNISTNetwork(),\n", 202 | " loss=nn.CrossEntropyLoss(),\n", 203 | " y_target=1,\n", 204 | " poisoned_rate=0.05,\n", 205 | " poisoned_transform_train_index=0,\n", 206 | " poisoned_transform_test_index=0,\n", 207 | " poisoned_target_transform_index=0,\n", 208 | " schedule=None,\n", 209 | " seed=global_seed,\n", 210 | " deterministic=deterministic,\n", 211 | " reflection_candidates = reflection_images,\n", 212 | ")\n", 213 | "\n", 214 | "poisoned_train_dataset, poisoned_test_dataset = refool.get_poisoned_dataset()\n", 215 | "print(\"============train_images============\")\n", 216 | "for i in range(10):\n", 217 | " show_img(poisoned_train_dataset[i][0])\n", 218 | "print(\"============test_images============\")\n", 219 | "for i in range(10):\n", 220 | " show_img(poisoned_test_dataset[i][0])" 221 | ] 222 | } 223 | ], 224 | "metadata": { 225 | "interpreter": { 226 | "hash": "0f3a347edc68809802ca80b72bd825e96565764d57572064874c38c4d7a65333" 227 | }, 228 | "kernelspec": { 229 | "display_name": "Python 3.8.12 64-bit ('tc18': conda)", 230 | "language": "python", 231 | "name": "python3" 232 | }, 233 | "language_info": { 234 | "codemirror_mode": { 235 | "name": "ipython", 236 | "version": 3 237 | }, 238 | "file_extension": ".py", 239 | "mimetype": "text/x-python", 240 | "name": "python", 241 | "nbconvert_exporter": "python", 242 | "pygments_lexer": "ipython3", 243 | "version": "3.8.12" 244 | }, 245 | "orig_nbformat": 4 246 | }, 247 | "nbformat": 4, 248 | "nbformat_minor": 2 249 | } 250 | -------------------------------------------------------------------------------- /tests/train_AutoEncoder.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Train an antoencoder network for antoencoder defense. 3 | ''' 4 | 5 | 6 | import torch 7 | import torchvision 8 | from torchvision.transforms import Compose, ToTensor, RandomHorizontalFlip 9 | 10 | import core 11 | 12 | 13 | # ========== Set global settings ========== 14 | global_seed = 666 15 | deterministic = True 16 | torch.manual_seed(global_seed) 17 | CUDA_VISIBLE_DEVICES = '1' 18 | datasets_root_dir = '../datasets' 19 | 20 | 21 | # ========== AutoEncoder3x32x32_CIFAR-10 ========== 22 | autoencoder_defense = core.AutoEncoderDefense( 23 | autoencoder=core.models.AutoEncoder(img_size=(3, 32, 32)), 24 | pretrain=None, 25 | seed=global_seed, 26 | deterministic=deterministic 27 | ) 28 | 29 | dataset = torchvision.datasets.CIFAR10 30 | transform_train = Compose([ 31 | # RandomHorizontalFlip(), 32 | ToTensor() 33 | ]) 34 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 35 | transform_test = Compose([ 36 | ToTensor() 37 | ]) 38 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 39 | 40 | # Train AutoEncoder Network (schedule is modified from https://github.com/chenjie/PyTorch-CIFAR-10-autoencoder/blob/master/main.py) 41 | schedule = { 42 | 'device': 'GPU', 43 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 44 | 'GPU_num': 1, 45 | 46 | 'batch_size': 16, 47 | 'num_workers': 2, 48 | 49 | 'lr': 0.001, 50 | 'betas': (0.9, 0.999), 51 | 'eps': 1e-08, 52 | 'weight_decay': 0, 53 | 'amsgrad': False, 54 | 55 | 'schedule': [], 56 | 'gamma': 0.1, 57 | 58 | 'epochs': 100, 59 | 60 | 'log_iteration_interval': 100, 61 | 'test_epoch_interval': 10, 62 | 'save_epoch_interval': 10, 63 | 64 | 'save_dir': 'experiments', 65 | 'experiment_name': 'AutoEncoder3x32x32_CIFAR-10_train' 66 | } 67 | 68 | autoencoder_defense.train_autoencoder(trainset, testset, schedule) 69 | 70 | 71 | 72 | 73 | -------------------------------------------------------------------------------- /tests/train_Benign.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the test code of Benign training. 3 | ''' 4 | 5 | 6 | import os.path as osp 7 | 8 | import cv2 9 | import torch 10 | import torch.nn as nn 11 | import torchvision 12 | from torchvision.datasets import DatasetFolder 13 | from torchvision.transforms import Compose, ToTensor, RandomHorizontalFlip, ToPILImage, Resize 14 | 15 | import core 16 | 17 | 18 | # ========== Set global settings ========== 19 | global_seed = 666 20 | deterministic = True 21 | torch.manual_seed(global_seed) 22 | CUDA_VISIBLE_DEVICES = '0' 23 | datasets_root_dir = '../datasets' 24 | 25 | 26 | # ========== BaselineMNISTNetwork_MNIST_Benign ========== 27 | dataset = torchvision.datasets.MNIST 28 | 29 | transform_train = Compose([ 30 | ToTensor() 31 | ]) 32 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 33 | 34 | transform_test = Compose([ 35 | ToTensor() 36 | ]) 37 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 38 | 39 | pattern = torch.zeros((28, 28), dtype=torch.uint8) 40 | pattern[-3:, -3:] = 255 41 | weight = torch.zeros((28, 28), dtype=torch.float32) 42 | weight[-3:, -3:] = 1.0 43 | 44 | badnets = core.BadNets( 45 | train_dataset=trainset, 46 | test_dataset=testset, 47 | model=core.models.BaselineMNISTNetwork(), 48 | loss=nn.CrossEntropyLoss(), 49 | y_target=1, 50 | poisoned_rate=0.05, 51 | pattern=pattern, 52 | weight=weight, 53 | seed=global_seed, 54 | deterministic=deterministic 55 | ) 56 | 57 | # Train Benign Model (schedule is set by yamengxi.) 58 | schedule = { 59 | 'device': 'GPU', 60 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 61 | 'GPU_num': 1, 62 | 63 | 'benign_training': True, # Train Benign Model 64 | 'batch_size': 128, 65 | 'num_workers': 2, 66 | 67 | 'lr': 0.1, 68 | 'momentum': 0.9, 69 | 'weight_decay': 5e-4, 70 | 'gamma': 0.1, 71 | 'schedule': [150, 180], 72 | 73 | 'epochs': 200, 74 | 75 | 'log_iteration_interval': 100, 76 | 'test_epoch_interval': 10, 77 | 'save_epoch_interval': 10, 78 | 79 | 'save_dir': 'experiments', 80 | 'experiment_name': 'BaselineMNISTNetwork_MNIST_Benign' 81 | } 82 | badnets.train(schedule) 83 | 84 | 85 | # ========== ResNet-18_CIFAR-10_Benign ========== 86 | dataset = torchvision.datasets.CIFAR10 87 | 88 | transform_train = Compose([ 89 | RandomHorizontalFlip(), 90 | ToTensor() 91 | ]) 92 | trainset = dataset(datasets_root_dir, train=True, transform=transform_train, download=True) 93 | 94 | transform_test = Compose([ 95 | ToTensor() 96 | ]) 97 | testset = dataset(datasets_root_dir, train=False, transform=transform_test, download=True) 98 | 99 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 100 | pattern[-3:, -3:] = 255 101 | weight = torch.zeros((32, 32), dtype=torch.float32) 102 | weight[-3:, -3:] = 1.0 103 | 104 | badnets = core.BadNets( 105 | train_dataset=trainset, 106 | test_dataset=testset, 107 | model=core.models.ResNet(18), 108 | loss=nn.CrossEntropyLoss(), 109 | y_target=1, 110 | poisoned_rate=0.05, 111 | pattern=pattern, 112 | weight=weight, 113 | seed=global_seed, 114 | deterministic=deterministic 115 | ) 116 | 117 | # Train Benign Model (schedule is the same as https://github.com/THUYimingLi/Open-sourced_Dataset_Protection/blob/main/CIFAR/train_standard.py) 118 | schedule = { 119 | 'device': 'GPU', 120 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 121 | 'GPU_num': 1, 122 | 123 | 'benign_training': True, # Train Benign Model 124 | 'batch_size': 128, 125 | 'num_workers': 2, 126 | 127 | 'lr': 0.1, 128 | 'momentum': 0.9, 129 | 'weight_decay': 5e-4, 130 | 'gamma': 0.1, 131 | 'schedule': [150, 180], 132 | 133 | 'epochs': 200, 134 | 135 | 'log_iteration_interval': 100, 136 | 'test_epoch_interval': 10, 137 | 'save_epoch_interval': 10, 138 | 139 | 'save_dir': 'experiments', 140 | 'experiment_name': 'ResNet-18_CIFAR-10_Benign' 141 | } 142 | badnets.train(schedule) 143 | 144 | 145 | # ========== ResNet-18_GTSRB_Benign ========== 146 | transform_train = Compose([ 147 | ToPILImage(), 148 | Resize((32, 32)), 149 | ToTensor() 150 | ]) 151 | trainset = DatasetFolder( 152 | root=osp.join(datasets_root_dir, 'GTSRB', 'train'), # please replace this with path to your training set 153 | loader=cv2.imread, 154 | extensions=('png',), 155 | transform=transform_train, 156 | target_transform=None, 157 | is_valid_file=None) 158 | 159 | transform_test = Compose([ 160 | ToPILImage(), 161 | Resize((32, 32)), 162 | ToTensor() 163 | ]) 164 | testset = DatasetFolder( 165 | root=osp.join(datasets_root_dir, 'GTSRB', 'testset'), # please replace this with path to your test set 166 | loader=cv2.imread, 167 | extensions=('png',), 168 | transform=transform_test, 169 | target_transform=None, 170 | is_valid_file=None) 171 | 172 | 173 | pattern = torch.zeros((32, 32), dtype=torch.uint8) 174 | pattern[-3:, -3:] = 255 175 | weight = torch.zeros((32, 32), dtype=torch.float32) 176 | weight[-3:, -3:] = 1.0 177 | 178 | badnets = core.BadNets( 179 | train_dataset=trainset, 180 | test_dataset=testset, 181 | model=core.models.ResNet(18, 43), 182 | loss=nn.CrossEntropyLoss(), 183 | y_target=1, 184 | poisoned_rate=0.05, 185 | pattern=pattern, 186 | weight=weight, 187 | poisoned_transform_train_index=2, 188 | poisoned_transform_test_index=2, 189 | seed=global_seed, 190 | deterministic=deterministic 191 | ) 192 | 193 | # Train Benign Model (schedule is the same as https://github.com/THUYimingLi/Open-sourced_Dataset_Protection/blob/main/GTSRB/train_standard.py) 194 | schedule = { 195 | 'device': 'GPU', 196 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 197 | 'GPU_num': 1, 198 | 199 | 'benign_training': True, # Train Benign Model 200 | 'batch_size': 128, 201 | 'num_workers': 2, 202 | 203 | 'lr': 0.01, 204 | 'momentum': 0.9, 205 | 'weight_decay': 5e-4, 206 | 'gamma': 0.1, 207 | 'schedule': [20], 208 | 209 | 'epochs': 30, 210 | 211 | 'log_iteration_interval': 100, 212 | 'test_epoch_interval': 10, 213 | 'save_epoch_interval': 10, 214 | 215 | 'save_dir': 'experiments', 216 | 'experiment_name': 'ResNet-18_GTSRB_Benign' 217 | } 218 | badnets.train(schedule) 219 | -------------------------------------------------------------------------------- /tests/train_REFINE.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Train an input transformation module for REFINE defense. 3 | ''' 4 | 5 | 6 | import sys 7 | import os 8 | sys.path.append(os.getcwd()) 9 | import torch 10 | import core 11 | import argparse 12 | from getmodel import GetModel 13 | from getdataset import GetDataset 14 | 15 | parser = argparse.ArgumentParser(description='PyTorch REFINE') 16 | parser.add_argument('--gpu_id', default='0', type=str, help='id(s) for CUDA_VISIBLE_DEVICES') 17 | parser.add_argument('--dataset', default='CIFAR10', type=str) 18 | parser.add_argument('--model', default='ResNet18', type=str) 19 | parser.add_argument('--attack', default='BadNets', type=str) 20 | parser.add_argument('--tlabel', default=0, type=int) 21 | 22 | args = parser.parse_args() 23 | state = {k: v for k, v in args._get_kwargs()} 24 | 25 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id 26 | 27 | # ========== Set training model ========== 28 | getmodel = GetModel(args.dataset, args.model, args.attack) 29 | my_model = getmodel.get_model() 30 | 31 | # ========== Set global settings ========== 32 | global_seed = 666 33 | deterministic = True 34 | torch.manual_seed(global_seed) 35 | CUDA_VISIBLE_DEVICES = args.gpu_id 36 | save_path = f'REFINE/{args.dataset}/{args.model}/{args.attack}' 37 | if not os.path.exists(save_path): 38 | os.makedirs(save_path) 39 | 40 | batch_size = 128 41 | num_workers = 4 42 | 43 | getdataset = GetDataset(args.dataset, args.model, args.attack, args.tlabel) 44 | trainset, testset = getdataset.get_benign_dataset() 45 | poisoned_trainset, poisoned_testset = getdataset.get_poisoned_dataset() 46 | 47 | if args.dataset == 'CIFAR10': 48 | size_map = 32 49 | first_channel = 64 50 | elif args.dataset == 'ImageNet50': 51 | size_map = 224 52 | first_channel = 32 53 | 54 | 55 | defense = core.REFINE( 56 | unet=core.models.UNetLittle(args=None, n_channels=3, n_classes=3, first_channels=first_channel), 57 | model=my_model, 58 | num_classes=10, 59 | seed=global_seed, 60 | deterministic=deterministic 61 | ) 62 | 63 | schedule = { 64 | 'device': 'GPU', 65 | 'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 66 | 'GPU_num': 1, 67 | 68 | # 'test_model': model_path, 69 | 'batch_size': batch_size, 70 | 'num_workers': num_workers, 71 | 72 | 'lr': 0.01, 73 | 'betas': (0.9, 0.999), 74 | 'eps': 1e-08, 75 | 'weight_decay': 0, 76 | 'amsgrad': False, 77 | 78 | 'schedule': [100, 130], 79 | 'gamma': 0.1, 80 | 81 | 'epochs': 150, 82 | 83 | 'log_iteration_interval': 100, 84 | 'test_epoch_interval': 10, 85 | 'save_epoch_interval': 10, 86 | 87 | 'save_dir': save_path, 88 | 'experiment_name': f'REFINE_train' 89 | } 90 | defense.train_unet(trainset, testset, schedule) 91 | 92 | unet_trained = defense.unet 93 | arr_shuffle_defined = defense.arr_shuffle --------------------------------------------------------------------------------