attack.py

import logging
from typing import Dict

import torch
from copy import deepcopy

from models.model import Model
from models.nc_model import NCModel
#from synthesizers.synthesizer import Synthesizer
from losses.loss_functions import compute_all_losses_and_grads
from utils.min_norm_solvers import MGDASolver
from utils.parameters import Params

logger = logging.getLogger('logger')


class Attack:
    params: Params
    #synthesizer: Synthesizer
    nc_model: Model
    nc_optim: torch.optim.Optimizer
    # fixed_model: Model

    def __init__(self, params):
        self.params = params
        #self.synthesizer = synthesizer

        # NC hyper params
        # if 'neural_cleanse' in self.params.loss_tasks:
        #     self.nc_model = NCModel(params.input_shape[1]).to(params.device)
        #     self.nc_optim = torch.optim.Adam(self.nc_model.parameters(), 0.01)

    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 = ['normal']
        #loss_tasks = self.params.loss_tasks.copy() if attack else ['normal']
        #batch_back = self.synthesizer.make_backdoor_batch(batch, attack=attack)
        scale = dict()

        # if 'neural_cleanse' in loss_tasks:
        #     self.neural_cleanse_part1(model, batch, batch_back)

        if len(loss_tasks) == 1:
            loss_values, grads = compute_all_losses_and_grads(
                loss_tasks,
                self, model, criterion, batch, 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

    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

    # def neural_cleanse_part1(self, model, batch, batch_back):
    #     self.nc_model.zero_grad()
    #     model.zero_grad()

    #     self.nc_model.switch_grads(True)
    #     model.switch_grads(False)
    #     output = model(self.nc_model(batch.inputs))
    #     nc_tasks = ['neural_cleanse_part1', 'mask_norm']

    #     criterion = torch.nn.CrossEntropyLoss(reduction='none')

    #     loss_values, grads = compute_all_losses_and_grads(nc_tasks,
    #                                                       self, model,
    #                                                       criterion, batch,
    #                                                       batch_back,
    #                                                       compute_grad=False
    #                                                       )
    #     # Using NC paper params
    #     logger.info(loss_values)
    #     loss = 0.999 * loss_values['neural_cleanse_part1'] + 0.001 * loss_values['mask_norm']
    #     loss.backward()
    #     self.nc_optim.step()

    #     self.nc_model.switch_grads(False)
    #     model.switch_grads(True)


    # def fl_scale_update(self, local_update: Dict[str, torch.Tensor]):
    #     for name, value in local_update.items():
    #         value.mul_(self.params.fl_weight_scale)