fl_trainer.py

import numpy as np

import torch
import torch.nn.functional as F

from utils import *
from defense import *

from models.vgg import get_vgg_model
import pandas as pd

from torch.nn.utils import parameters_to_vector, vector_to_parameters
import datasets


class Net(nn.Module):
    def __init__(self, num_classes):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        output = self.fc2(x)
        # output = F.log_softmax(x, dim=1)
        return output


def get_results_filename(poison_type, attack_method, model_replacement, project_frequency, defense_method, norm_bound,
                         prox_attack, fixed_pool=False, model_arch="vgg9"):
    filename = "{}_{}_{}".format(poison_type, model_arch, attack_method)
    if fixed_pool:
        filename += "_fixed_pool"

    if model_replacement:
        filename += "_with_replacement"
    else:
        filename += "_without_replacement"

    if attack_method == "pgd":
        filename += "_1_{}".format(project_frequency)

    if prox_attack:
        filename += "_prox_attack"

    if defense_method in ("norm-clipping", "norm-clipping-adaptive", "weak-dp"):
        filename += "_{}_m_{}".format(defense_method, norm_bound)
    elif defense_method in ("multi-metrics", "krum", "multi-krum", "rfa"):
        filename += "_{}".format(defense_method)

    filename += "_acc_results.csv"

    return filename


def calc_norm_diff(gs_model, vanilla_model, epoch, fl_round, mode="bad"):
    norm_diff = 0
    for p_index, p in enumerate(gs_model.parameters()):
        norm_diff += torch.norm(list(gs_model.parameters())[p_index] - list(vanilla_model.parameters())[p_index]) ** 2
    norm_diff = torch.sqrt(norm_diff).item()
    if mode == "bad":
        # pdb.set_trace()
        logger.info(
            "===> ND `|w_bad-w_g|` in local epoch: {} | FL round: {} |, is {}".format(epoch, fl_round, norm_diff))
    elif mode == "normal":
        logger.info(
            "===> ND `|w_normal-w_g|` in local epoch: {} | FL round: {} |, is {}".format(epoch, fl_round, norm_diff))
    elif mode == "avg":
        logger.info(
            "===> ND `|w_avg-w_g|` in local epoch: {} | FL round: {} |, is {}".format(epoch, fl_round, norm_diff))

    return norm_diff


def fed_avg_aggregator(net_list, net_freq, device, model="lenet"):
    # net_avg = VGG('VGG11').to(device)
    if model == "lenet":
        net_avg = Net(num_classes=10).to(device)
    elif model in ("vgg9", "vgg11", "vgg13", "vgg16"):
        net_avg = get_vgg_model(model).to(device)
    whole_aggregator = []

    for p_index, p in enumerate(net_list[0].parameters()):
        # initial 
        params_aggregator = torch.zeros(p.size()).to(device)
        for net_index, net in enumerate(net_list):
            # we assume the adv model always comes to the beginning
            params_aggregator = params_aggregator + net_freq[net_index] * list(net.parameters())[p_index].data
        whole_aggregator.append(params_aggregator)

    for param_index, p in enumerate(net_avg.parameters()):
        p.data = whole_aggregator[param_index]
    return net_avg


def estimate_wg(model, device, train_loader, optimizer, epoch, log_interval, criterion):
    logger.info("Prox-attack: Estimating wg_hat")
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    if batch_idx % log_interval == 0:
        logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            epoch, batch_idx * len(data), len(train_loader.dataset),
                   100. * batch_idx / len(train_loader), loss.item()))


def train(model, device, train_loader, optimizer, epoch, log_interval, criterion, pgd_attack=False, eps=5e-4,
          model_original=None,
          proj="l_2", project_frequency=1, adv_optimizer=None, prox_attack=False, wg_hat=None):
    """
        train function for both honest nodes and adversary.
        NOTE: this trains only for one epoch
    """
    model.train()
    # get learning rate
    for param_group in optimizer.param_groups:
        eta = param_group['lr']

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        if pgd_attack:
            adv_optimizer.zero_grad()
        output = model(data)
        # loss = F.nll_loss(output, target)
        loss = criterion(output, target.long())
        if prox_attack:
            wg_hat_vec = parameters_to_vector(list(wg_hat.parameters()))
            model_vec = parameters_to_vector(list(model.parameters()))
            prox_term = torch.norm(wg_hat_vec - model_vec) ** 2
            loss = loss + prox_term

        loss.backward()
        if not pgd_attack:
            optimizer.step()
        else:
            if proj == "l_inf":
                w = list(model.parameters())
                n_layers = len(w)
                # adversarial learning rate
                eta = 0.001
                for i in range(len(w)):
                    # uncomment below line to restrict proj to some layers
                    if True:  # i == 6 or i == 8 or i == 10 or i == 0 or i == 18:
                        w[i].data = w[i].data - eta * w[i].grad.data
                        # projection step
                        m1 = torch.lt(torch.sub(w[i], model_original[i]), -eps)
                        m2 = torch.gt(torch.sub(w[i], model_original[i]), eps)
                        w1 = (model_original[i] - eps) * m1
                        w2 = (model_original[i] + eps) * m2
                        w3 = (w[i]) * (~(m1 + m2))
                        wf = w1 + w2 + w3
                        w[i].data = wf.data
            else:
                # do l2_projection
                adv_optimizer.step()
                w = list(model.parameters())
                w_vec = parameters_to_vector(w)
                model_original_vec = parameters_to_vector(model_original)
                # make sure you project on last iteration otherwise, high LR pushes you really far
                if (batch_idx % project_frequency == 0 or batch_idx == len(train_loader) - 1) and (
                        torch.norm(w_vec - model_original_vec) > eps):
                    # project back into norm ball
                    w_proj_vec = eps * (w_vec - model_original_vec) / torch.norm(
                        w_vec - model_original_vec) + model_original_vec
                    # plug w_proj back into model
                    vector_to_parameters(w_proj_vec, w)
                # for i in range(n_layers):
                #    # uncomment below line to restrict proj to some layers
                #    if True:#i == 16 or i == 17:
                #        w[i].data = w[i].data - eta * w[i].grad.data
                #        if torch.norm(w[i] - model_original[i]) > eps/n_layers:
                #            # project back to norm ball
                #            w_proj= (eps/n_layers)*(w[i]-model_original[i])/torch.norm(
                #                w[i]-model_original[i]) + model_original[i]
                #            w[i].data = w_proj

        if batch_idx % log_interval == 0:
            logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(model, device, test_loader, test_batch_size, criterion, mode="raw-task", dataset="cifar10",
         poison_type="fashion"):
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))

    if dataset in ("mnist", "emnist"):
        target_class = 7
        if mode == "raw-task":
            classes = [str(i) for i in range(10)]
        elif mode == "targetted-task":
            if poison_type == 'ardis':
                classes = [str(i) for i in range(10)]
            else:
                classes = ["T-shirt/top",
                           "Trouser",
                           "Pullover",
                           "Dress",
                           "Coat",
                           "Sandal",
                           "Shirt",
                           "Sneaker",
                           "Bag",
                           "Ankle boot"]
    elif dataset == "cifar10":
        classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
        # target_class = 2 for greencar, 9 for southwest
        if poison_type in ("howto", "greencar-neo"):
            target_class = 2
        else:
            target_class = 9

    model.eval()
    test_loss = 0
    correct = 0
    backdoor_correct = 0
    backdoor_tot = 0
    final_acc = 0
    task_acc = None

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output, 1)
            c = (predicted == target).squeeze()

            # test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            test_loss += criterion(output, target.long()).item()
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
            # check backdoor accuracy
            if poison_type == 'ardis':
                backdoor_index = torch.where(target == target_class)
                target_backdoor = torch.ones_like(target[backdoor_index])
                predicted_backdoor = predicted[backdoor_index]
                backdoor_correct += (predicted_backdoor == target_backdoor).sum().item()
                backdoor_tot = backdoor_index[0].shape[0]
                # logger.info("Target: {}".format(target_backdoor))
                # logger.info("Predicted: {}".format(predicted_backdoor))

            # for image_index in range(test_batch_size):
            for image_index in range(len(target)):
                label = target[image_index]
                class_correct[label] += c[image_index].item()
                class_total[label] += 1
    test_loss /= len(test_loader.dataset)

    if mode == "raw-task":
        for i in range(10):
            logger.info('Accuracy of %5s : %.2f %%' % (
                classes[i], 100 * class_correct[i] / class_total[i]))

            if i == target_class:
                task_acc = 100 * class_correct[i] / class_total[i]

        logger.info('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))
        final_acc = 100. * correct / len(test_loader.dataset)

    elif mode == "targetted-task":

        if dataset in ("mnist", "emnist"):
            for i in range(10):
                logger.info('Accuracy of %5s : %.2f %%' % (
                    classes[i], 100 * class_correct[i] / class_total[i]))
            if poison_type == 'ardis':
                # ensure 7 is being classified as 1
                logger.info('Backdoor Accuracy of %.2f : %.2f %%' % (
                    target_class, 100 * backdoor_correct / backdoor_tot))
                final_acc = 100 * backdoor_correct / backdoor_tot
            else:
                # trouser acc
                final_acc = 100 * class_correct[1] / class_total[1]

        elif dataset == "cifar10":
            logger.info('#### Targetted Accuracy of %5s : %.2f %%' % (
            classes[target_class], 100 * class_correct[target_class] / class_total[target_class]))
            final_acc = 100 * class_correct[target_class] / class_total[target_class]
    return final_acc, task_acc


class FederatedLearningTrainer:
    def __init__(self, *args, **kwargs):
        self.hyper_params = None

    def run(self, client_model, *args, **kwargs):
        raise NotImplementedError()


class FrequencyFederatedLearningTrainer(FederatedLearningTrainer):
    def __init__(self, arguments=None, *args, **kwargs):
        # self.poisoned_emnist_dataset = arguments['poisoned_emnist_dataset']
        self.vanilla_model = arguments['vanilla_model']
        self.net_avg = arguments['net_avg']
        self.net_dataidx_map = arguments['net_dataidx_map']
        self.num_nets = arguments['num_nets']
        self.part_nets_per_round = arguments['part_nets_per_round']
        self.fl_round = arguments['fl_round']
        self.local_training_period = arguments['local_training_period']
        self.adversarial_local_training_period = arguments['adversarial_local_training_period']
        self.args_lr = arguments['args_lr']
        self.args_gamma = arguments['args_gamma']
        self.attacking_fl_rounds = arguments['attacking_fl_rounds']
        self.poisoned_emnist_train_loader = arguments['poisoned_emnist_train_loader']
        self.clean_train_loader = arguments['clean_train_loader']
        self.vanilla_emnist_test_loader = arguments['vanilla_emnist_test_loader']
        self.targetted_task_test_loader = arguments['targetted_task_test_loader']
        self.batch_size = arguments['batch_size']
        self.test_batch_size = arguments['test_batch_size']
        self.log_interval = arguments['log_interval']
        self.device = arguments['device']
        self.num_dps_poisoned_dataset = arguments['num_dps_poisoned_dataset']
        self.defense_technique = arguments["defense_technique"]
        self.norm_bound = arguments["norm_bound"]
        self.attack_method = arguments["attack_method"]
        self.dataset = arguments["dataset"]
        self.model = arguments["model"]
        self.criterion = nn.CrossEntropyLoss()
        self.eps = arguments['eps']
        self.poison_type = arguments['poison_type']
        self.model_replacement = arguments['model_replacement']
        self.project_frequency = arguments['project_frequency']
        self.adv_lr = arguments['adv_lr']
        self.prox_attack = arguments['prox_attack']
        self.attack_case = arguments['attack_case']
        self.stddev = arguments['stddev']

        logger.info("Posion type! {}".format(self.poison_type))

        if self.poison_type == 'ardis':
            self.ardis_dataset = datasets.get_ardis_dataset()
            # exclude first 66 points because they are part of the adversary
            if self.attack_case == 'normal-case':
                self.ardis_dataset.data = self.ardis_dataset.data[66:]
            elif self.attack_case == 'almost-edge-case':
                self.ardis_dataset.data = self.ardis_dataset.data[66:132]
        elif self.poison_type == 'southwest':
            self.ardis_dataset = datasets.get_southwest_dataset(attack_case=self.attack_case)
        else:
            self.ardis_dataset = None

        if self.attack_method == "pgd":
            self.pgd_attack = True
        else:
            self.pgd_attack = False

        if arguments["defense_technique"] == "no-defense":
            self._defender = None
        elif arguments["defense_technique"] == "norm-clipping" or arguments[
            "defense_technique"] == "norm-clipping-adaptive":
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "weak-dp":
            # doesn't really add noise. just clips
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "multi-metrics":
            self._defender = Multi_metrics(num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "krum":
            self._defender = Krum(mode='krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "multi-krum":
            self._defender = Krum(mode='multi-krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "rfa":
            self._defender = RFA()
        elif arguments["defense_technique"] == "foolsgold":
            self._defender = FoolsGold()
        else:
            NotImplementedError("Unsupported defense method !")

    def run(self):
        main_task_acc = []
        raw_task_acc = []
        backdoor_task_acc = []
        fl_iter_list = []
        adv_norm_diff_list = []
        wg_norm_list = []
        epoch_choice_list = []
        global_choice_list = []
        # let's conduct multi-round training
        for flr in range(1, self.fl_round + 1):
            logger.info("##### attack fl rounds: {}".format(self.attacking_fl_rounds))
            g_user_indices = []

            if self.defense_technique == "norm-clipping-adaptive":
                # experimental
                norm_diff_collector = []

            if flr in self.attacking_fl_rounds:
                # randomly select participating clients
                # in this current version, we sample `part_nets_per_round-1` per FL round since we assume attacker will always participates
                selected_node_indices = np.random.choice(self.num_nets, size=self.part_nets_per_round - 1,
                                                         replace=False)
                num_data_points = [len(self.net_dataidx_map[i]) for i in selected_node_indices]
                total_num_dps_per_round = sum(num_data_points) + self.num_dps_poisoned_dataset

                logger.info("FL round: {}, total num data points: {}, num dps poisoned: {}".format(flr, num_data_points,
                                                                                                   self.num_dps_poisoned_dataset))

                net_freq = [self.num_dps_poisoned_dataset / total_num_dps_per_round] + [
                    num_data_points[i] / total_num_dps_per_round for i in range(self.part_nets_per_round - 1)]
                logger.info("Net freq: {}, FL round: {} with adversary".format(net_freq, flr))
                # pdb.set_trace()

                # we need to reconstruct the net list at the beginning
                net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
                logger.info("################## Starting fl round: {}".format(flr))

                model_original = list(self.net_avg.parameters())
                # super hacky but I'm doing this for the prox-attack
                wg_clone = copy.deepcopy(self.net_avg)
                wg_hat = None
                v0 = torch.nn.utils.parameters_to_vector(model_original)
                wg_norm_list.append(torch.norm(v0).item())

                # start the FL process
                for net_idx, net in enumerate(net_list):
                    # net  = net_list[net_idx]
                    if net_idx == 0:
                        global_user_idx = -1  # we assign "-1" as the indices of the attacker in global user indices
                        pass
                    else:
                        global_user_idx = selected_node_indices[net_idx - 1]
                        dataidxs = self.net_dataidx_map[global_user_idx]
                        if self.attack_case == "edge-case":
                            train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size,
                                                               self.test_batch_size,
                                                               dataidxs)  # also get the data loader
                        elif self.attack_case in ("normal-case", "almost-edge-case"):
                            train_dl_local, _ = get_dataloader_normal_case(self.dataset, './data', self.batch_size,
                                                                           self.test_batch_size, dataidxs,
                                                                           user_id=global_user_idx,
                                                                           num_total_users=self.num_nets,
                                                                           poison_type=self.poison_type,
                                                                           ardis_dataset=self.ardis_dataset,
                                                                           attack_case=self.attack_case)  # also get the data loader
                        else:
                            NotImplementedError("Unsupported attack case ...")

                    g_user_indices.append(global_user_idx)
                    if net_idx == 0:
                        logger.info("@@@@@@@@ Working on client: {}, which is Attacker".format(net_idx))
                    else:
                        logger.info(
                            "@@@@@@@@ Working on client: {}, which is Global user: {}".format(net_idx, global_user_idx))

                    criterion = nn.CrossEntropyLoss()
                    optimizer = optim.SGD(net.parameters(), lr=self.args_lr * self.args_gamma ** (flr - 1),
                                          momentum=0.9,
                                          weight_decay=1e-4)  # epoch, net, train_loader, optimizer, criterion
                    adv_optimizer = optim.SGD(net.parameters(), lr=self.adv_lr * self.args_gamma ** (flr - 1),
                                              momentum=0.9,
                                              weight_decay=1e-4)  # looks like adversary needs same lr to hide with others
                    prox_optimizer = optim.SGD(wg_clone.parameters(), lr=self.args_lr * self.args_gamma ** (flr - 1),
                                               momentum=0.9, weight_decay=1e-4)
                    for param_group in optimizer.param_groups:
                        logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))

                    if net_idx == 0:
                        if self.prox_attack:
                            # estimate w_hat
                            for inner_epoch in range(1, self.local_training_period + 1):
                                estimate_wg(wg_clone, self.device, self.clean_train_loader, prox_optimizer, inner_epoch,
                                            log_interval=self.log_interval, criterion=self.criterion)
                            wg_hat = wg_clone

                        for e in range(1, self.adversarial_local_training_period + 1):
                            # we always assume net index 0 is adversary
                            if self.defense_technique in ('krum', 'multi-krum', 'multi-metrics'):
                                train(net, self.device, self.poisoned_emnist_train_loader, optimizer, e,
                                      log_interval=self.log_interval, criterion=self.criterion,
                                      pgd_attack=self.pgd_attack, eps=self.eps * self.args_gamma ** (flr - 1),
                                      model_original=model_original, project_frequency=self.project_frequency,
                                      adv_optimizer=adv_optimizer,
                                      prox_attack=self.prox_attack, wg_hat=wg_hat)
                            else:
                                train(net, self.device, self.poisoned_emnist_train_loader, optimizer, e,
                                      log_interval=self.log_interval, criterion=self.criterion,
                                      pgd_attack=self.pgd_attack, eps=self.eps, model_original=model_original,
                                      project_frequency=self.project_frequency, adv_optimizer=adv_optimizer,
                                      prox_attack=self.prox_attack, wg_hat=wg_hat)

                        test(net, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size,
                             criterion=self.criterion, mode="raw-task", dataset=self.dataset,
                             poison_type=self.poison_type)
                        test(net, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size,
                             criterion=self.criterion, mode="targetted-task", dataset=self.dataset,
                             poison_type=self.poison_type)

                        # if model_replacement scale models
                        if self.model_replacement:
                            v = torch.nn.utils.parameters_to_vector(net.parameters())
                            logger.info("Attacker before scaling : Norm = {}".format(torch.norm(v)))
                            # adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")
                            # logger.info("||w_bad - w_avg|| before scaling = {}".format(adv_norm_diff))

                            for idx, param in enumerate(net.parameters()):
                                param.data = (param.data - model_original[idx]) * (
                                            total_num_dps_per_round / self.num_dps_poisoned_dataset) + model_original[
                                                 idx]
                            v = torch.nn.utils.parameters_to_vector(net.parameters())
                            logger.info("Attacker after scaling : Norm = {}".format(torch.norm(v)))

                        # at here we can check the distance between w_bad and w_g i.e. `\|w_bad - w_g\|_2`
                        # we can print the norm diff out for debugging
                        adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr,
                                                       mode="bad")
                        adv_norm_diff_list.append(adv_norm_diff)

                        if self.defense_technique == "norm-clipping-adaptive":
                            # experimental
                            norm_diff_collector.append(adv_norm_diff)
                    else:
                        for e in range(1, self.local_training_period + 1):
                            train(net, self.device, train_dl_local, optimizer, e, log_interval=self.log_interval,
                                  criterion=self.criterion)
                            # at here we can check the distance between w_normal and w_g i.e. `\|w_bad - w_g\|_2`
                        # we can print the norm diff out for debugging
                        honest_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e,
                                                          fl_round=flr, mode="normal")

                        if self.defense_technique == "norm-clipping-adaptive":
                            # experimental
                            norm_diff_collector.append(honest_norm_diff)

            else:
                selected_node_indices = np.random.choice(self.num_nets, size=self.part_nets_per_round, replace=False)
                num_data_points = [len(self.net_dataidx_map[i]) for i in selected_node_indices]
                total_num_dps_per_round = sum(num_data_points)

                net_freq = [num_data_points[i] / total_num_dps_per_round for i in range(self.part_nets_per_round)]
                logger.info("Net freq: {}, FL round: {} without adversary".format(net_freq, flr))

                # we need to reconstruct the net list at the beginning
                net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
                logger.info("################## Starting fl round: {}".format(flr))

                # start the FL process
                for net_idx, net in enumerate(net_list):
                    global_user_idx = selected_node_indices[net_idx]
                    dataidxs = self.net_dataidx_map[global_user_idx]

                    if self.attack_case == "edge-case":
                        train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size,
                                                           self.test_batch_size, dataidxs)  # also get the data loader
                    elif self.attack_case in ("normal-case", "almost-edge-case"):
                        train_dl_local, _ = get_dataloader_normal_case(self.dataset, './data', self.batch_size,
                                                                       self.test_batch_size, dataidxs,
                                                                       user_id=global_user_idx,
                                                                       num_total_users=self.num_nets,
                                                                       poison_type=self.poison_type,
                                                                       ardis_dataset=self.ardis_dataset,
                                                                       attack_case=self.attack_case)  # also get the data loader
                    else:
                        NotImplementedError("Unsupported attack case ...")

                    g_user_indices.append(global_user_idx)

                    logger.info(
                        "@@@@@@@@ Working on client: {}, which is Global user: {}".format(net_idx, global_user_idx))

                    criterion = nn.CrossEntropyLoss()
                    optimizer = optim.SGD(net.parameters(), lr=self.args_lr * self.args_gamma ** (flr - 1),
                                          momentum=0.9,
                                          weight_decay=1e-4)  # epoch, net, train_loader, optimizer, criterion
                    for param_group in optimizer.param_groups:
                        logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))

                    for e in range(1, self.local_training_period + 1):
                        train(net, self.device, train_dl_local, optimizer, e, log_interval=self.log_interval,
                              criterion=self.criterion)

                    honest_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr,
                                                      mode="normal")

                    if self.defense_technique == "norm-clipping-adaptive":
                        # experimental
                        norm_diff_collector.append(honest_norm_diff)

                adv_norm_diff_list.append(0)
                model_original = list(self.net_avg.parameters())
                v0 = torch.nn.utils.parameters_to_vector(model_original)
                wg_norm_list.append(torch.norm(v0).item())
            attack_round = False
            if flr in self.attacking_fl_rounds:
                attack_round = True
            ### conduct defense here:
            if self.defense_technique == "no-defense":
                pass
            elif self.defense_technique == "norm-clipping":
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "norm-clipping-adaptive":
                # we will need to adapt the norm diff first before the norm diff clipping
                logger.info("#### Let's Look at the Nom Diff Collector : {} ....; Mean: {}".format(norm_diff_collector,
                                                                                                   np.mean(
                                                                                                       norm_diff_collector)))
                self._defender.norm_bound = np.mean(norm_diff_collector)
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "weak-dp":
                # this guy is just going to clip norm. No noise added here
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net,
                                        global_model=self.net_avg, )
            elif self.defense_technique == "multi-metrics":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                         num_dps=[self.num_dps_poisoned_dataset] + num_data_points,
                                                         g_user_indices=g_user_indices,
                                                         device=self.device,
                                                         attack_round=attack_round)
            elif self.defense_technique == "krum":
                net_list, net_freq, epoch_choice, global_choice = self._defender.exec(client_models=net_list,
                                                                                      num_dps=[
                                                                                                  self.num_dps_poisoned_dataset] + num_data_points,
                                                                                      g_user_indices=g_user_indices,
                                                                                      device=self.device,
                                                                                      attack_round=attack_round)
            elif self.defense_technique == "multi-krum":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                         num_dps=[self.num_dps_poisoned_dataset] + num_data_points,
                                                         g_user_indices=g_user_indices,
                                                         device=self.device,
                                                         attack_round=attack_round)
            elif self.defense_technique == "rfa":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                         net_freq=net_freq,
                                                         maxiter=500,
                                                         eps=1e-5,
                                                         ftol=1e-7,
                                                         device=self.device)
            elif self.defense_technique == "foolsgold":
                net_list, net_freq = self._defender.exec(client_models=net_list, net_freq=net_freq,
                                                         names=selected_node_indices, device=self.device)
            else:
                NotImplementedError("Unsupported defense method !")

            # after local training periods
            self.net_avg = fed_avg_aggregator(net_list, net_freq, device=self.device, model=self.model)
            if self.defense_technique == "weak-dp":
                # add noise to self.net_avg
                noise_adder = AddNoise(stddev=self.stddev)
                noise_adder.exec(client_model=self.net_avg,
                                 device=self.device)

            v = torch.nn.utils.parameters_to_vector(self.net_avg.parameters())
            logger.info("############ Averaged Model : Norm {}".format(torch.norm(v)))

            calc_norm_diff(gs_model=self.net_avg, vanilla_model=self.vanilla_model, epoch=0, fl_round=flr, mode="avg")

            logger.info("Measuring the accuracy of the averaged global model, FL round: {} ...".format(flr))

            overall_acc, raw_acc = test(self.net_avg, self.device, self.vanilla_emnist_test_loader,
                                        test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task",
                                        dataset=self.dataset, poison_type=self.poison_type)
            backdoor_acc, _ = test(self.net_avg, self.device, self.targetted_task_test_loader,
                                   test_batch_size=self.test_batch_size, criterion=self.criterion,
                                   mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)

            fl_iter_list.append(flr)
            main_task_acc.append(overall_acc)
            raw_task_acc.append(raw_acc)
            backdoor_task_acc.append(backdoor_acc)

        df = pd.DataFrame({'fl_iter': fl_iter_list,
                           'main_task_acc': main_task_acc,
                           'backdoor_acc': backdoor_task_acc,
                           'raw_task_acc': raw_task_acc,
                           'adv_norm_diff': adv_norm_diff_list,
                           'wg_norm': wg_norm_list
                           })

        if self.poison_type == 'ardis':
            # add a row showing initial accuracies
            df1 = pd.DataFrame(
                {'fl_iter': [0], 'main_task_acc': [88], 'backdoor_acc': [11], 'raw_task_acc': [0], 'adv_norm_diff': [0],
                 'wg_norm': [0]})
            df = pd.concat([df1, df])

        results_filename = get_results_filename(self.poison_type, self.attack_method, self.model_replacement,
                                                self.project_frequency,
                                                self.defense_technique, self.norm_bound, self.prox_attack, False,
                                                self.model)

        df.to_csv(results_filename, index=False)
        logger.info("Wrote accuracy results to: {}".format(results_filename))

        # save model net_avg
        # torch.save(self.net_avg.state_dict(), "./checkpoint/emnist_lenet_10epoch.pt")


class FixedPoolFederatedLearningTrainer(FederatedLearningTrainer):
    def __init__(self, arguments=None, *args, **kwargs):

        # self.poisoned_emnist_dataset = arguments['poisoned_emnist_dataset']
        self.vanilla_model = arguments['vanilla_model']
        self.net_avg = arguments['net_avg']
        self.net_dataidx_map = arguments['net_dataidx_map']
        self.num_nets = arguments['num_nets']
        self.part_nets_per_round = arguments['part_nets_per_round']
        self.fl_round = arguments['fl_round']
        self.local_training_period = arguments['local_training_period']
        self.adversarial_local_training_period = arguments['adversarial_local_training_period']
        self.args_lr = arguments['args_lr']
        self.args_gamma = arguments['args_gamma']
        self.attacker_pool_size = arguments['attacker_pool_size']
        self.poisoned_emnist_train_loader = arguments['poisoned_emnist_train_loader']
        self.clean_train_loader = arguments['clean_train_loader']
        self.vanilla_emnist_test_loader = arguments['vanilla_emnist_test_loader']
        self.targetted_task_test_loader = arguments['targetted_task_test_loader']
        self.batch_size = arguments['batch_size']
        self.test_batch_size = arguments['test_batch_size']
        self.log_interval = arguments['log_interval']
        self.device = arguments['device']
        self.dataset = arguments["dataset"]
        self.model = arguments["model"]
        self.num_dps_poisoned_dataset = arguments['num_dps_poisoned_dataset']
        self.defense_technique = arguments["defense_technique"]
        self.norm_bound = arguments["norm_bound"]
        self.attack_method = arguments["attack_method"]
        self.criterion = nn.CrossEntropyLoss()
        self.eps = arguments['eps']
        self.dataset = arguments["dataset"]
        self.poison_type = arguments['poison_type']
        self.model_replacement = arguments['model_replacement']
        self.project_frequency = arguments['project_frequency']
        self.adv_lr = arguments['adv_lr']
        self.prox_attack = arguments['prox_attack']
        self.attack_case = arguments['attack_case']
        self.stddev = arguments['stddev']

        logger.info("Posion type! {}".format(self.poison_type))

        if self.attack_method == "pgd":
            self.pgd_attack = True
        else:
            self.pgd_attack = False

        if self.poison_type == 'ardis':
            self.ardis_dataset = datasets.get_ardis_dataset()
            # exclude first 66 points because they are part of the adversary
            if self.attack_case == 'normal-case':
                self.ardis_dataset.data = self.ardis_dataset.data[66:]
            elif self.attack_case == 'almost-edge-case':
                self.ardis_dataset.data = self.ardis_dataset.data[66:132]
        elif self.poison_type == 'southwest':
            self.ardis_dataset = datasets.get_southwest_dataset(attack_case=self.attack_case)
        else:
            self.ardis_dataset = None

        if arguments["defense_technique"] == "no-defense":
            self._defender = None
        elif arguments["defense_technique"] == "norm-clipping" or arguments[
            "defense_technique"] == "norm-clipping-adaptive":
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "weak-dp":
            # doesn't really add noise. just clips
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "multi-metrics":
            self._defender = Multi_metrics(num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "krum":
            self._defender = Krum(mode='krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "multi-krum":
            self._defender = Krum(mode='multi-krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "rfa":
            self._defender = RFA()
        elif arguments["defense_technique"] == "foolsgold":
            self._defender = FoolsGold()
        else:
            NotImplementedError("Unsupported defense method !")

        self.attacker_pool = np.random.choice(self.num_nets, self.attacker_pool_size, replace=False)
        # arguments["attacker_pool"] = self.attacker_pool

    def run(self, arguments=None, *args, **kwargs):
        main_task_acc = []
        raw_task_acc = []
        backdoor_task_acc = []
        fl_iter_list = []
        adv_norm_diff_list = []
        wg_norm_list = []
        epoch_choice_list = []
        global_choice_list = []
        # let's conduct multi-round training
        for flr in range(1, self.fl_round + 1):
            # randomly select participating clients
            # in this current version, we sample `part_nets_per_round` per FL round since we assume attacker will always participates
            g_user_indices = []
            selected_node_indices = np.random.choice(self.num_nets, size=self.part_nets_per_round, replace=False)
            arguments["selected_node_indices"] = selected_node_indices

            selected_attackers = [idx for idx in selected_node_indices if idx in self.attacker_pool]
            arguments["selected_attackers"] = selected_attackers
            selected_honest_users = [idx for idx in selected_node_indices if idx not in self.attacker_pool]
            logger.info("Selected Attackers in FL iteration-{}: {}".format(flr, selected_attackers))
            arguments["selected_honest_users"] = selected_honest_users

            num_data_points = []
            for sni in selected_node_indices:
                if sni in selected_attackers:
                    num_data_points.append(self.num_dps_poisoned_dataset)
                else:
                    num_data_points.append(len(self.net_dataidx_map[sni]))

            total_num_dps_per_round = sum(num_data_points)
            net_freq = [num_data_points[i] / total_num_dps_per_round for i in range(self.part_nets_per_round)]
            logger.info("Net freq: {}, FL round: {} with adversary".format(net_freq, flr))

            # we need to reconstruct the net list at the beginning
            net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
            logger.info("################## Starting fl round: {}".format(flr))
            model_original = list(self.net_avg.parameters())
            # super hacky but I'm doing this for the prox-attack
            wg_clone = copy.deepcopy(self.net_avg)
            wg_hat = None
            v0 = torch.nn.utils.parameters_to_vector(model_original)
            wg_norm_list.append(torch.norm(v0).item())

            #     # start the FL process
            for net_idx, global_user_idx in enumerate(selected_node_indices):
                g_user_indices.append(global_user_idx)
                net = net_list[net_idx]
                if global_user_idx in selected_attackers:
                    pass
                else:
                    dataidxs = self.net_dataidx_map[global_user_idx]
                    # train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size,
                    #                                self.test_batch_size, dataidxs) # also get the data loader

                    # add p-percent edge-case attack here:
                    if self.attack_case == "edge-case":
                        train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size,
                                                           self.test_batch_size, dataidxs)  # also get the data loader
                    elif self.attack_case in ("normal-case", "almost-edge-case"):
                        train_dl_local, _ = get_dataloader_normal_case(self.dataset, './data', self.batch_size,
                                                                       self.test_batch_size, dataidxs,
                                                                       user_id=global_user_idx,
                                                                       num_total_users=self.num_nets,
                                                                       poison_type=self.poison_type,
                                                                       ardis_dataset=self.ardis_dataset,
                                                                       attack_case=self.attack_case)  # also get the data loader
                    else:
                        NotImplementedError("Unsupported attack case ...")

                logger.info(
                    "@@@@@@@@ Working on client (global-index): {}, which {}-th user in the current round".format(
                        global_user_idx, net_idx))

                # criterion = nn.CrossEntropyLoss()
                # optimizer = optim.SGD(net.parameters(), lr=self.args_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4) # epoch, net, train_loader, optimizer, criterion
                # for param_group in optimizer.param_groups:
                #    logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))                

                criterion = nn.CrossEntropyLoss()
                optimizer = optim.SGD(net.parameters(), lr=self.args_lr * self.args_gamma ** (flr - 1), momentum=0.9,
                                      weight_decay=1e-4)  # epoch, net, train_loader, optimizer, criterion
                adv_optimizer = optim.SGD(net.parameters(), lr=self.adv_lr * self.args_gamma ** (flr - 1), momentum=0.9,
                                          weight_decay=1e-4)  # looks like adversary needs same lr to hide with others
                prox_optimizer = optim.SGD(wg_clone.parameters(), lr=self.args_lr * self.args_gamma ** (flr - 1),
                                           momentum=0.9, weight_decay=1e-4)
                for param_group in optimizer.param_groups:
                    logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))

                current_adv_norm_diff_list = []
                if global_user_idx in selected_attackers:
                    # for e in range(1, self.adversarial_local_training_period+1):
                    #    # we always assume net index 0 is adversary
                    #    train(net, self.device, self.poisoned_emnist_train_loader, optimizer, e, log_interval=self.log_interval, criterion=self.criterion)

                    # logger.info("=====> Measuring the model performance of the poisoned model after attack ....")
                    # test(net, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
                    # test(net, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)
                    # # at here we can check the distance between w_bad and w_g i.e. `\|w_bad - w_g\|_2`
                    # calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")

                    if self.prox_attack:
                        # estimate w_hat
                        for inner_epoch in range(1, self.local_training_period + 1):
                            estimate_wg(wg_clone, self.device, self.clean_train_loader, prox_optimizer, inner_epoch,
                                        log_interval=self.log_interval, criterion=self.criterion)
                        wg_hat = wg_clone

                    for e in range(1, self.adversarial_local_training_period + 1):
                        # we always assume net index 0 is adversary
                        if self.defense_technique in ('krum', 'multi-krum', 'multi-metrics'):
                            train(net, self.device, self.poisoned_emnist_train_loader, optimizer, e,
                                  log_interval=self.log_interval, criterion=self.criterion,
                                  pgd_attack=self.pgd_attack, eps=self.eps * self.args_gamma ** (flr - 1),
                                  model_original=model_original, project_frequency=self.project_frequency,
                                  adv_optimizer=adv_optimizer,
                                  prox_attack=self.prox_attack, wg_hat=wg_hat)
                        else:
                            train(net, self.device, self.poisoned_emnist_train_loader, optimizer, e,
                                  log_interval=self.log_interval, criterion=self.criterion,
                                  pgd_attack=self.pgd_attack, eps=self.eps, model_original=model_original,
                                  project_frequency=self.project_frequency, adv_optimizer=adv_optimizer,
                                  prox_attack=self.prox_attack, wg_hat=wg_hat)

                    test(net, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size,
                         criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
                    test(net, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size,
                         criterion=self.criterion, mode="targetted-task", dataset=self.dataset,
                         poison_type=self.poison_type)

                    # if model_replacement scale models
                    if self.model_replacement:
                        v = torch.nn.utils.parameters_to_vector(net.parameters())
                        logger.info("Attacker before scaling : Norm = {}".format(torch.norm(v)))
                        # adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")
                        # logger.info("||w_bad - w_avg|| before scaling = {}".format(adv_norm_diff))

                        for idx, param in enumerate(net.parameters()):
                            param.data = (param.data - model_original[idx]) * (
                                        total_num_dps_per_round / self.num_dps_poisoned_dataset) + model_original[idx]
                        v = torch.nn.utils.parameters_to_vector(net.parameters())
                        logger.info("Attacker after scaling : Norm = {}".format(torch.norm(v)))

                    # at here we can check the distance between w_bad and w_g i.e. `\|w_bad - w_g\|_2`
                    # we can print the norm diff out for debugging
                    adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr,
                                                   mode="bad")
                    current_adv_norm_diff_list.append(adv_norm_diff)

                    if self.defense_technique == "norm-clipping-adaptive":
                        # experimental
                        norm_diff_collector.append(adv_norm_diff)
                else:
                    # for e in range(1, self.local_training_period+1):
                    #    train(net, self.device, train_dl_local, optimizer, e, log_interval=self.log_interval, criterion=self.criterion)                
                    # # at here we can check the distance between w_normal and w_g i.e. `\|w_bad - w_g\|_2`
                    # calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="normal")

                    for e in range(1, self.local_training_period + 1):
                        train(net, self.device, train_dl_local, optimizer, e, log_interval=self.log_interval,
                              criterion=self.criterion)
                        # at here we can check the distance between w_normal and w_g i.e. `\|w_bad - w_g\|_2`
                    # we can print the norm diff out for debugging
                    honest_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr,
                                                      mode="normal")

                    if self.defense_technique == "norm-clipping-adaptive":
                        # experimental
                        norm_diff_collector.append(honest_norm_diff)

            ### conduct defense here:
            if self.defense_technique == "no-defense":
                pass
            elif self.defense_technique == "norm-clipping":
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "norm-clipping-adaptive":
                # we will need to adapt the norm diff first before the norm diff clipping
                logger.info("#### Let's Look at the Nom Diff Collector : {} ....; Mean: {}".format(norm_diff_collector,
                                                                                                   np.mean(
                                                                                                       norm_diff_collector)))
                self._defender.norm_bound = np.mean(norm_diff_collector)
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "weak-dp":
                # this guy is just going to clip norm. No noise added here
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net,
                                        global_model=self.net_avg, )
            elif self.defense_technique == "multi-metrics":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                                                      num_dps=[self.num_dps_poisoned_dataset] + num_data_points,
                                                                                      g_user_indices=g_user_indices,
                                                                                      device=self.device)
            elif self.defense_technique == "krum":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                                                      num_dps=[self.num_dps_poisoned_dataset] + num_data_points,
                                                                                      g_user_indices=g_user_indices,
                                                                                      device=self.device)
            elif self.defense_technique == "multi-krum":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                                                      num_dps=[self.num_dps_poisoned_dataset] + num_data_points,
                                                                                      g_user_indices=g_user_indices,
                                                                                      device=self.device)
            elif self.defense_technique == "rfa":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                         net_freq=net_freq,
                                                         maxiter=500,
                                                         eps=1e-5,
                                                         ftol=1e-7,
                                                         device=self.device)
            elif self.defense_technique == "foolsgold":
                net_list, net_freq = self._defender.exec(client_models=net_list, names=selected_node_indices,
                                                         device=self.device)
            else:
                NotImplementedError("Unsupported defense method !")

            # after local training periods
            self.net_avg = fed_avg_aggregator(net_list, net_freq, device=self.device, model=self.model)

            if self.defense_technique == "weak-dp":
                # add noise to self.net_avg
                noise_adder = AddNoise(stddev=self.stddev)
                noise_adder.exec(client_model=self.net_avg,
                                 device=self.device)

            calc_norm_diff(gs_model=self.net_avg, vanilla_model=self.net_avg, epoch=0, fl_round=flr, mode="avg")

            logger.info("Measuring the accuracy of the averaged global model, FL round: {} ...".format(flr))
            # overall_acc = test(self.net_avg, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
            # backdoor_acc = test(self.net_avg, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)
            overall_acc, raw_acc = test(self.net_avg, self.device, self.vanilla_emnist_test_loader,
                                        test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task",
                                        dataset=self.dataset, poison_type=self.poison_type)
            backdoor_acc, _ = test(self.net_avg, self.device, self.targetted_task_test_loader,
                                   test_batch_size=self.test_batch_size, criterion=self.criterion,
                                   mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)

            fl_iter_list.append(flr)
            main_task_acc.append(overall_acc)
            raw_task_acc.append(raw_acc)
            backdoor_task_acc.append(backdoor_acc)
            if len(current_adv_norm_diff_list) == 0:
                adv_norm_diff_list.append(0)
            else:
                # if you have multiple adversaries in a round, average their norm diff
                adv_norm_diff_list.append(1.0 * sum(current_adv_norm_diff_list) / len(current_adv_norm_diff_list))

        df = pd.DataFrame({'fl_iter': fl_iter_list,
                           'main_task_acc': main_task_acc,
                           'backdoor_acc': backdoor_task_acc,
                           'raw_task_acc': raw_task_acc,
                           'adv_norm_diff': adv_norm_diff_list,
                           'wg_norm': wg_norm_list
                           })

        if self.poison_type == 'ardis':
            # add a row showing initial accuracies
            df1 = pd.DataFrame(
                {'fl_iter': [0], 'main_task_acc': [88], 'backdoor_acc': [11], 'raw_task_acc': [0], 'adv_norm_diff': [0],
                 'wg_norm': [0]})
            df = pd.concat([df1, df])

        results_filename = get_results_filename(self.poison_type, self.attack_method, self.model_replacement,
                                                self.project_frequency,
                                                self.defense_technique, self.norm_bound, self.prox_attack,
                                                fixed_pool=True, model_arch=self.model)
        df.to_csv(results_filename, index=False)
        logger.info("Wrote accuracy results to: {}".format(results_filename))