联邦学习后门攻击代码阅读——backdoors101

最新推荐文章于 2024-10-06 17:37:54 发布

Mr. Water

最新推荐文章于 2024-10-06 17:37:54 发布

阅读量3.6k

点赞数 12

分类专栏： Python 联邦学习

本文链接：https://blog.csdn.net/narcissus2_/article/details/113972161

版权

Python 同时被 2 个专栏收录

3 篇文章 0 订阅

订阅专栏

联邦学习

2 篇文章 0 订阅

订阅专栏

代码地址

https://github.com/MrWater98/backdoors101

论文地址

https://arxiv.org/abs/1807.00459
https://arxiv.org/abs/2005.03823

代码阅读目标

清楚概括代码运行过程。
能将代码和论文相对应。
收集目前还不了解的内容。
只阅读和联邦学习相关的内容

代码结构

│ attack.py
│ helper.py
│ requirements.txt
│ training.py
├─configs
　　│ cifar10_params.yaml
　　│ cifar_fed.yaml
　　│ imagenet_params.yaml
　　│ mnist_params.yaml
　　│ pipa_params.yaml
├─dataset
　　│ celeba.py
　　│ multi_mnist_loader.py
　　│ pipa.py
　　│ vggface.py
├─losses
　　│ loss_functions.py
├─metrics
　　│ accuracy_metric.py
　　│ metric.py
　　│ test_loss_metric.py
├─models
　　│ face_ident.py
　　│ model.py
　　│ resnet.py
　　│ simple.py
　　│ vgg.py
　　│ word_model.py
　　│ init.py
├─src
　　│ attack_vectors.png
　　│ calculator.png
　　│ complex.png
　　│ pipa.png
　　│ pixel_vs_semantic.png
├─synthesizers
　　│ complex_synthesizer.py
　　│ pattern_synthesizer.py
　　│ physical_synthesizer.py
　　│ singlepixel_synthesizer.py
　　│ synthesizer.py
├─tasks
　　│ │ batch.py
　　│ │ celeba_helper.py
　　│ │ cifar10_task.py
　　│ │ imagenet_task.py
　　│ │ imdb_helper.py
　　│ │ mnist_task.py
　　│ │ multimnist_helper.py
　　│ │ pipa_task.py
　　│ │ task.py
　　│ │ vggface_helper.py
　　│ │
　　│ └─fl
　　　　│ cifarfed_task.py
　　　　│ fl_emnist_task.py
　　　　│ fl_reddit_task.py
　　　　│ fl_task.py
　　　　│ fl_user.py
　　└─utils
　　│ │ index.html
　　│ │ min_norm_solvers.py
　　│ │ parameters.py
　　│ │ utils.py
　　│ │ init.py

运行方法

安装所有依赖pip install -r requirements.txt
创建两个目录runs和saved_models
通过tensorboard --logdir=runs/创建tensorboard
通过python training.py --name mnist --params configs/mnist_params.yaml --commit none来运行其中的一个训练集
1. 这里运行的是mnist_params.yaml
2. 如果需要运行的是联邦学习的，则需要调整name后面的参数和params后面的参数。

梳理流程

1. 先观察trainning.py

主要的工作就是根据获得的param来构成Helper，从而帮助建立整一个框架。
然后run才负责运行整个框架。

if __name__ == '__main__':
    # 将所有的命令行参数都解读到parser
    parser = argparse.ArgumentParser(description='Backdoors')
    parser.add_argument('--params', dest='params', default='utils/params.yaml')
    parser.add_argument('--name', dest='name', required=True)
    parser.add_argument('--commit', dest='commit',
                        default=get_current_git_hash())

    args = parser.parse_args()
    # 第二个参数设定的.yaml最重要，name只是确认你创建的文件的名字
    with open(args.params) as f:
        params = yaml.load(f, Loader=yaml.FullLoader)
    
    params['current_time'] = datetime.now().strftime('%b.%d_%H.%M.%S')
    params['commit'] = args.commit
    params['name'] = args.name
    # Helper读取params
    helper = Helper(params)
    logger.warning(create_table(params))
    
    try:
        # 参数fl来自于cifar_fed.yaml
        if helper.params.fl:
            fl_run(helper)
        else:
            run(helper)

2. Helper.py

Helper的工作包括要：
- 判断是否是联邦学习
- 生成Synthesizer，为攻击模型准备。
- attack包括了Synthesizer和一些计算loss的函数，可以进行多任务的操作。
- tb_writer是tensorboard可视化结果的一个工具。

class Helper:
    params: Params = None
    # https://docs.python.org/3/library/typing.html#typing.Union
    # 这应该是意为这要不是Task，要不是FederatedLearningTask
    task: Union[Task, FederatedLearningTask] = None
    # 来源于 https://github.com/MrWater98/backdoors101 
    # 将未backdoored的输入转化为backdoored的输入
    synthesizer: Synthesizer = None
    # 包括了多个人物的同步器和损失率的计算
    attack: Attack = None
    # tensorboard的结果
    tb_writer: SummaryWriter = None

    def __init__(self, params):
        self.params = Params(**params)

        self.times = {'backward': list(), 'forward': list(), 'step': list(),
                      'scales': list(), 'total': list(), 'poison': list()}
        if self.params.random_seed is not None:
            self.fix_random(self.params.random_seed)
        # 创建结果的文件夹
        self.make_folders()
        # 找到训练用的xx_task文件，用默认构造函数获取构造后的结果
        self.make_task()
        # 基本同上
        self.make_synthesizer()
        # 拿到Attakc对象
        self.attack = Attack(self.params, self.synthesizer)
        # neural cleanse 识别和减轻神经网络中的后门攻击的手段
        self.nc = True if 'neural_cleanse' in self.params.loss_tasks else False

        self.best_loss = float('inf')

2.1. make_folder()

主要功能

根据params的log是否为True判断是否要在params_folder_path创建文件夹。
文件夹中还包含run.html，一些画图内容包含在里面。
还会根据tb是否为True来判断是否要使用Tensorboard作图。

2.2. make_task()

主要功能
- 判断使用的的数据集以及需要调用数据集的位置。
- 通过获取module的名字来引入这个类以及其内部的一些方法。
- self.task通过默认构造函数以及helper的params获取task

2.2.1. fl_task.py

本函数从make_task()中进入
主要的目有:
- 创建一个训练好的残差网络/恢复一个残差网络。
- 判断使用gpu还是cpu进行训练。
- 选择使用的评价标准（这里使用的是cross entrophy）

class FederatedLearningTask(Task):
    fl_train_loaders: List[Any] = None
    ignored_weights = ['num_batches_tracked']#['tracked', 'running']
    adversaries: List[int] = None

    def init_task(self):
        # load_data对应的是cifarfed_task.py
        self.load_data()
        # build_model就只是创建一个18层的残差网络
        self.model = self.build_model()
        # resume_model是当有训练过的模型后，则会使用之前的模型
        self.resume_model()
        # 使用cpu或者gpu训练，这里主要是使用cpu
        self.model = self.model.to(self.params.device)

        self.local_model = self.build_model().to(self.params.device)
        # 直接用entrophy了
        self.criterion = self.make_criterion()
        # 这个是选择攻击的样本
        self.adversaries = self.sample_adversaries()
		# 评价矩阵
        self.metrics = [AccuracyMetric(), TestLossMetric(self.criterion)]
        
        self.set_input_shape()
        return

2.2.1.1. load_data()<-cifarfed_task.py

cifared_task.py中的类CifarFedTask本身继承自FederatedLearningTask和Cifar10Task。
所以load_cifar_data实际上来源于cifar10_task.py
(https://blog.csdn.net/u010165147/article/details/78633858)

def load_data(self) -> None:

        self.load_cifar_data()
        if self.params.fl_sample_dirichlet:
            # 使用狄利克雷分布来进行用户训练数据的取样
            indices_per_participant = self.sample_dirichlet_train_data(
                self.params.fl_total_participants,
                alpha=self.params.fl_dirichlet_alpha)
            train_loaders = [(pos, self.get_train(indices)) for pos, indices in
                             indices_per_participant.items()]
        else:
            # 否则就平均分即可
            all_range = list(range(len(self.train_dataset)))
            random.shuffle(all_range)
            train_loaders = [self.get_train_old(all_range, pos)
                             for pos in range(self.params.fl_total_participants)]
        self.fl_train_loaders = train_loaders
        return

load_cifar_data()本身不难理解，我已经把注释写在下面了：
- 随机裁剪的原因参见：[随机裁剪的原因]

def load_cifar_data(self):
        # 如果参数中要做transform，那么则
        if self.params.transform_train:
            transform_train = transforms.Compose([
                # 做一个随机裁剪
                transforms.RandomCrop(32, padding=4),
                # 做一个随机的上下翻转
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                # 做一个正则化
                self.normalize,
            ])
        else:
            transform_train = transforms.Compose([
                transforms.ToTensor(),
                self.normalize,
            ])
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            self.normalize,
        ])
        # 下载一下训练的数据
        self.train_dataset = torchvision.datasets.CIFAR10(
            root=self.params.data_path,
            train=True,
            download=True,
            transform=transform_train)
        # 如果已经有要下毒的照片，则不是用semantic backdoor
        if self.params.poison_images:
            self.train_loader = self.remove_semantic_backdoors()
        else:
            self.train_loader = DataLoader(self.train_dataset,
                                           batch_size=self.params.batch_size,
                                           shuffle=True,
                                           num_workers=0)
        # 下载一下测试集的数据
        self.test_dataset = torchvision.datasets.CIFAR10(
            root=self.params.data_path,
            train=False,
            download=True,
            transform=transform_test)
        self.test_loader = DataLoader(self.test_dataset,
                                      batch_size=self.params.test_batch_size,
                                      shuffle=False, num_workers=0)
        # 定义一下类
        self.classes = ('plane', 'car', 'bird', 'cat',
                        'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
        return True

2.2.1.2. build_model

函数来自于cifar10_task.py，主要目的是创建一个残差网络，输出看有多少个类。

# 创建残差网络
    def build_model(self) -> nn.Module:
        if self.params.pretrained:
            model = resnet18(pretrained=True)
            model.fc = nn.Linear(512, len(self.classes))
        else:
            model = resnet18(pretrained=False,
                                  num_classes=len(self.classes))
        return model

2.2.1.3. sample_adversaries

主要的目的就是获取攻击的用户下标ID，通过不同的采样手段。

def sample_adversaries(self) -> List[int]:
        adversaries_ids = []
        # 对应cifar_fed.yaml第45行，
        if self.params.fl_number_of_adversaries == 0:
            # vanilla 寻常的，没有新意的
            logger.warning(f'Running vanilla FL, no attack.')

        elif self.params.fl_single_epoch_attack is None:
            adversaries_ids = random.sample(
                range(self.params.fl_total_participants),
                self.params.fl_number_of_adversaries)
            logger.warning(f'Attacking over multiple epochs with following '
                           f'users compromised: {adversaries_ids}.')
        else:
            logger.warning(f'Attack only on epoch: '
                           f'{self.params.fl_single_epoch_attack} with '
                           f'{self.params.fl_number_of_adversaries} compromised'
                           f' users.')
         return adversaries_ids

2.3. fl_run和run_fl_round

我们现在已经获得好了所有需要的材料，进入跑模型的阶段了。

def fl_run(hlpr: Helper):
    for epoch in range(hlpr.params.start_epoch,
                       hlpr.params.epochs + 1):
        run_fl_round(hlpr, epoch)
        metric = test(hlpr, epoch, backdoor=False)
        test(hlpr, epoch, backdoor=True)

        hlpr.save_model(hlpr.task.model, epoch, metric)


def run_fl_round(hlpr, epoch):
    # 获得global的模型
    global_model = hlpr.task.model
    # 获得local的模型
    local_model = hlpr.task.local_model

    round_participants = hlpr.task.sample_users_for_round(epoch)
    weight_accumulator = hlpr.task.get_empty_accumulator()
    # tqdm是python的进度条库，基本是基于对象迭代
    for user in tqdm(round_participants):
        # 将参数从global_model复制到local_model
        hlpr.task.copy_params(global_model, local_model)
        # 一个对象，会保存当前状态，并根据梯度更新参数
        optimizer = hlpr.task.make_optimizer(local_model)
        for local_epoch in range(hlpr.params.fl_local_epochs):
            # 如果是恶意的用户，则执行进攻的训练
            if user.compromised:
                train(hlpr, local_epoch, local_model, optimizer,
                      user.train_loader, attack=True)
            # 如果是非恶意的用户，则执行非进攻的训练
            else:
                train(hlpr, local_epoch, local_model, optimizer,
                      user.train_loader, attack=False)
        # 然后来更新global的模型
        local_update = hlpr.task.get_fl_update(local_model, global_model)
        # 如果用户是恶意用户，还会更新梯度
        if user.compromised:
            hlpr.attack.fl_scale_update(local_update)
        # 存疑，感觉是积累当前的权重变化
        hlpr.task.accumulate_weights(weight_accumulator, local_update)
    # 所有用户完成之后，更新全局的模型
    hlpr.task.update_global_model(weight_accumulator, global_model)

这里fl_run都比较好理解，但是run_fl_round相对来说比较难理解。我把注释写在了代码上。

2.3.1. copy_params

def copy_params(self, global_model, local_model):
        # 复制global模型中的全部参数
        local_state = local_model.state_dict()
        for name, param in global_model.state_dict().items():
            if name in local_state and name not in self.ignored_weights:
                local_state[name].copy_(param)

2.3.2. make_optimizer

def make_optimizer(self, model=None) -> Optimizer:
        if model is None:
            model = self.model
        # 随机梯度下降
        if self.params.optimizer == 'SGD':
            optimizer = optim.SGD(model.parameters(),
                                  lr=self.params.lr,
                                  weight_decay=self.params.decay,
                                  momentum=self.params.momentum)
        # 动量和自适应学习率优化下降，SGD升级版
        elif self.params.optimizer == 'Adam':
            optimizer = optim.Adam(model.parameters(),
                                   lr=self.params.lr,
                                   weight_decay=self.params.decay)
        else:
            raise ValueError(f'No optimizer: {self.optimizer}')

        return optimizer

2.4. train

现在进入了训练的阶段，可以观察代码是如何训练和攻击当前模型的了。

def train(hlpr: Helper, epoch, model, optimizer, train_loader, attack=True):
    # criterion指的是评价指标
    criterion = hlpr.task.criterion
    model.train()
    for i, data in enumerate(train_loader):
        batch = hlpr.task.get_batch(i, data)
        # 把梯度设置成0，在计算反向传播的时候一般都会这么操作，原因未知
        model.zero_grad()
        # 主要进行攻击的代码
        # 可以看blind backdoor xxx
        loss = hlpr.attack.compute_blind_loss(model, criterion, batch, attack)
        loss.backward()
        # 使用optimizer.step()之后，模型才会更新
        optimizer.step()
        # 打印的函数
        hlpr.report_training_losses_scales(i, epoch)
        if i == hlpr.params.max_batch_id:
            break
    return

论文解读

理解整一个攻击的逻辑最主要的是要理解攻击的含义以及背后代表的数学意义是什么。
简单来说，后门攻击的含义就是要使得某一类图片被误分类。例如我想要打广告，就让大量不相关的图片都识别成我的品牌名。
实际操作的方式就是给模型错误的标签，让模型错误地被分类即可。
数学含义就是最小化这个误分类的任务loss函数，实际上这和训练一个正常的网络区别并不大。
上面这种情况主要是在secure aggregation的情况下，但如果模型引入了anomaly detection，那么我们还需要增加一个loss任务，即最小化被认为是anomaly user的loss。

2.4.1. get_batch

主要目的就是获得当前的batch，关于epoch，batch的相关解释可以看：epoch, batch, iteration的相关解释

2.4.2. compute_blind_loss(self, model, criterion, batch, attack)

可以说是整个项目最关键的函数，主要的目的就是计算loss用于反向传播。

def compute_blind_loss(self, model, criterion, batch, attack):
        """
        :param model:
        :param criterion:
        :param batch:
        :param attack: Do not attack at all. Ignore all the parameters
        :return:
        """
        batch = batch.clip(self.params.clip_batch)
        loss_tasks = self.params.loss_tasks.copy() if attack else ['normal']
        batch_back = self.synthesizer.make_backdoor_batch(batch, attack=attack)
        scale = dict()
        if len(loss_tasks) == 1:
            loss_values, grads = compute_all_losses_and_grads(
                loss_tasks,
                self, model, criterion, batch, batch_back, compute_grad=False
            )
        elif self.params.loss_balance == 'MGDA':

            loss_values, grads = compute_all_losses_and_grads(
                loss_tasks,
                self, model, criterion, batch, batch_back, compute_grad=True)

            if len(loss_tasks) > 1:
                scale = MGDASolver.get_scales(grads, loss_values,
                                              self.params.mgda_normalize,
                                              loss_tasks)
        elif self.params.loss_balance == 'fixed':
            loss_values, grads = compute_all_losses_and_grads(
                loss_tasks,
                self, model, criterion, batch, batch_back, compute_grad=False)

            for t in loss_tasks:
                scale[t] = self.params.fixed_scales[t]
        else:
            raise ValueError(f'Please choose between `MGDA` and `fixed`.')

        if len(loss_tasks) == 1:
            scale = {loss_tasks[0]: 1.0}

        blind_loss = self.scale_losses(loss_tasks, loss_values, scale)

        return blind_loss

2.4.2.1. make_backdoor_batch(batch, test, attack)

函数来自于Synthesizer.py
主要的功能是实施后门以及注入后门的动作。

def make_backdoor_batch(self, batch: Batch, test=False, attack=True) -> Batch:

        # Don't attack if only normal loss task.
        if (not attack) or (self.params.loss_tasks == ['normal'] and not test):
            return batch

        if test:
            attack_portion = batch.batch_size
        else:
            # 攻击的的位置来源于从数据集中随机取样
            attack_portion = round(
                batch.batch_size * self.params.poisoning_proportion)

        backdoored_batch = batch.clone()
        # 传入batch和portion
        self.apply_backdoor(backdoored_batch, attack_portion)

        return backdoored_batch
        
def apply_backdoor(self, batch, attack_portion):
        """
        Modifies only a portion of the batch (represents batch poisoning).

        :param batch:
        :return:
        """
        
        self.synthesize_inputs(batch=batch, attack_portion=attack_portion)
        self.synthesize_labels(batch=batch, attack_portion=attack_portion)

        return

2.4.2.2. synthesize_inputs和synthesize_labels

主要的功能是在图片中加入pattern，以及给他们分配专门的label。
这里的label是8，对应的是ship。
实施的pattern是随机生成的，并且在Synthesizer已经生成好了。
注意，这里生成的都是instance，而并不是实施到了模型中。

def synthesize_inputs(self, batch, attack_portion=None):
        # 在attakch的portion上加入pattern
        pattern, mask = self.get_pattern()
        batch.inputs[:attack_portion] = (1 - mask) * \
                                        batch.inputs[:attack_portion] + \
                                        mask * pattern

        return

    def synthesize_labels(self, batch, attack_portion=None):
        # 在attack的portion上加入backdoor_label，这里backdoor_label是8
        batch.labels[:attack_portion].fill_(self.params.backdoor_label)

        return

2.4.2.3. compute_all_loses_and_grads

我们的loss_balance是MGDA，这个的意思是Multiple-gradient descent algorithm，它用于当我们有多个优化目标的时候，这个多任务优化的原理我还不是很了解，但是我发现它需要两个模型的grads,loss_values。
我们当前在配置文件中只有两个任务，注入后门和其他的普通任务。所以实际上的攻击任务就转化成了多目标优化问题。

loss_tasks:
  - backdoor
  - normal
#  - nc_adv
#  - ewc
#  - latent
#  - latent_fixed

下面的函数主要目标是获得loss_values和grads通过compute_xxx_loss。
compute_grad用于了解是否要计算梯度，也决定了上面的函数是否会返回梯度。

def compute_all_losses_and_grads(loss_tasks, attack, model, criterion,
                                 batch, batch_back,
                                 compute_grad=None):
    grads = {}
    loss_values = {}
    for t in attack.params.loss_tasks:
        # if compute_grad:
        #     model.zero_grad()
        if t == 'normal':
            loss_values[t], grads[t] = compute_normal_loss(attack.params,
                                                           model,
                                                           criterion,
                                                           batch.inputs,
                                                           batch.labels,
                                                           grads=compute_grad)
        elif t == 'backdoor':
            loss_values[t], grads[t] = compute_backdoor_loss(attack.params,
                                                             model,
                                                             criterion,
                                                             batch_back.inputs,
                                                             batch_back.labels,
                                                             grads=compute_grad)
        elif t == 'spectral_evasion':
            loss_values[t], grads[t] = compute_spectral_evasion_loss(
                attack.params,
                model,
                attack.fixed_model,
                batch.inputs,
                grads=compute_grad)
        elif t == 'sentinet_evasion':
            loss_values[t], grads[t] = compute_sentinet_evasion(
                attack.params,
                model,
                batch.inputs,
                batch_back.inputs,
                batch_back.labels,
                grads=compute_grad)
        elif t == 'mask_norm':
            loss_values[t], grads[t] = norm_loss(attack.params, model,
                                                 grads=compute_grad)
        elif t == 'nc':
            loss_values[t], grads[t] = compute_normal_loss(attack.params,
                                                           model,
                                                           criterion,
                                                           batch.inputs,
                                                           batch_back.labels,
                                                           grads=compute_grad,
                                                           )

        # if loss_values[t].mean().item() == 0.0:
        #     loss_values.pop(t)
        #     grads.pop(t)
        #     loss_tasks.remove(t)
    return loss_values, grads

2.4.2.4. compute_xxx_loss(params, model, criterion, inputs,labels, grads):

基本就是把输入丢到模型后，用criterion算一下loss。
梯度就从parameter里面把梯度拎出来。

def compute_normal_loss(params, model, criterion, inputs,
                        labels, grads):
    t = time.perf_counter()
    outputs = model(inputs)
    record_time(params, t, 'forward')
    loss = criterion(outputs, labels)

    if not params.dp:
        loss = loss.mean()

    if grads:
        t = time.perf_counter()
        grads = list(torch.autograd.grad(loss.mean(),
                                         [x for x in model.parameters() if
                                          x.requires_grad],
                                         retain_graph=True))
        record_time(params, t, 'backward')

    return loss, grads

def compute_backdoor_loss(params, model, criterion, inputs_back,
                          labels_back, grads=None):
    t = time.perf_counter()
    outputs = model(inputs_back)
    record_time(params, t, 'forward')

    if params.task == 'pipa':
        loss = criterion(outputs, labels_back)
        loss[labels_back == 0] *= 0.001
        if labels_back.sum().item() == 0.0:
            loss[:] = 0.0
        loss = loss.mean()
    else:
        loss = criterion(outputs, labels_back)
    if not params.dp:
        loss = loss.mean()

    if grads:
        grads = get_grads(params, model, loss)

    return loss, grads

2.4.2.5. scale_losses(self, loss_tasks, loss_values, scale)

对两个loss进行一个放缩的相加，放缩的比例来自于上面MGDASolver中的get_scales

def scale_losses(self, loss_tasks, loss_values, scale):
        blind_loss = 0
        for it, t in enumerate(loss_tasks):
            self.params.running_losses[t].append(loss_values[t].item())
            self.params.running_scales[t].append(scale[t])
            if it == 0:
                blind_loss = scale[t] * loss_values[t]
            else:
                blind_loss += scale[t] * loss_values[t]
        self.params.running_losses['total'].append(blind_loss.item())
        return blind_loss