jupyter_ver.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#args\n",
    "import argparse\n",
    "    \n",
    "parser = argparse.ArgumentParser()\n",
    "# federated arguments\n",
    "parser.add_argument('--epochs', type=int, default=50, help=\"rounds of training\")\n",
    "parser.add_argument('--num_users', type=int, default=50, help=\"number of users: K\")\n",
    "parser.add_argument('--ratio', type=float, default=1,  help=\"portion of iid user: \")\n",
    "parser.add_argument('--frac', type=float, default=0.2, help=\"the fraction of clients: C\")\n",
    "parser.add_argument('--local_ep', type=int, default=50, help=\"the number of local epochs: E\")\n",
    "parser.add_argument('--local_bs', type=int, default=128, help=\"local batch size: B\")\n",
    "parser.add_argument('--bs', type=int, default=128, help=\"test batch size\")\n",
    "parser.add_argument('--lr', type=float, default=0.01, help=\"learning rate\")\n",
    "parser.add_argument('--momentum', type=float, default=0.5, help=\"SGD momentum (default: 0.5)\")\n",
    "parser.add_argument('--split', type=str, default='user', help=\"train-test split type, user or sample\")\n",
    "\n",
    "# model arguments\n",
    "parser.add_argument('--model', type=str, default='resnet', help='model name')\n",
    "parser.add_argument('--kernel_num', type=int, default=9, help='number of each kind of kernel')\n",
    "parser.add_argument('--kernel_sizes', type=str, default='3,4,5',\n",
    "                    help='comma-separated kernel size to use for convolution')\n",
    "parser.add_argument('--norm', type=str, default='batch_norm', help=\"batch_norm, layer_norm, or None\")\n",
    "parser.add_argument('--num_filters', type=int, default=32, help=\"number of filters for conv nets\")\n",
    "parser.add_argument('--max_pool', type=str, default='True',\n",
    "                    help=\"Whether use max pooling rather than strided convolutions\")\n",
    "\n",
    "# other arguments\n",
    "parser.add_argument('--dataset', type=str, default='cifar', help=\"name of dataset,cifar,mnist\")\n",
    "parser.add_argument(\"--sample\", type=int, default=500, help=\"number of samples for each node\")\n",
    "parser.add_argument('--pattern', type=str,  default='iid', help='iid,noniid,1-9,iid-q')\n",
    "parser.add_argument('--num_classes', type=int, default=10, help=\"number of classes\")\n",
    "parser.add_argument('--num_channels', type=int, default=1, help=\"number of channels of imges\")\n",
    "parser.add_argument('--gpu', type=int, default=0, help=\"GPU ID, -1 for CPU\")\n",
    "parser.add_argument('--stopping_rounds', type=int, default=10, help='rounds of early stopping')\n",
    "parser.add_argument('--verbose', action='store_true', help='verbose print')\n",
    "parser.add_argument('--seed', type=int, default=1, help='random seed (default: 1)')\n",
    "parser.add_argument('--all_clients', action='store_true', help='aggregation over all clients')\n",
    "parser.add_argument('--alg',type=str,default='avg',help='avg or fed')\n",
    "args = parser.parse_args([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sample\n",
    "import numpy as np\n",
    "from torchvision import datasets, transforms\n",
    "import torch\n",
    "# from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torch.utils.data import random_split\n",
    "def Dataset_config(dataset, num_users, pattern):\n",
    "    if dataset == 'mnist':\n",
    "        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n",
    "        dataset_train = datasets.MNIST(root='./data/', train=True, download=True,\n",
    "                                       transform=trans_mnist)\n",
    "        dataset_test = datasets.MNIST(root='./data/', train=False, download=True,\n",
    "                                      transform=trans_mnist)\n",
    "        X_train, y_train = dataset_train.data, dataset_train.targets\n",
    "        X_train = X_train.data.numpy()\n",
    "        y_train = y_train.data.numpy()\n",
    "\n",
    "    elif dataset == 'cifar':\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.RandomCrop(32, padding=4),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "        ])\n",
    "\n",
    "        transform_test = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "        ])\n",
    "\n",
    "\n",
    "        dataset_train = datasets.CIFAR10('./data/', train=True, download=True, transform=transform_train)\n",
    "        dataset_test = datasets.CIFAR10('./data/', train=False, download=True, transform=transform_test)\n",
    "                \n",
    "        X_train, y_train = dataset_train.data, dataset_train.targets\n",
    "        y_train = np.array(y_train)\n",
    "        \n",
    "    elif dataset =='fashion':\n",
    "        transformations = transforms.Compose([transforms.ToTensor(),])\n",
    "        dataset_train = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transformations)\n",
    "        dataset_test = datasets.FashionMNIST('./data', download=True, train=False, transform=transformations)\n",
    "        X_train, y_train = dataset_train.data, dataset_train.targets\n",
    "                     \n",
    "        \n",
    "    else:\n",
    "        exit('Error: unrecognized dataset')\n",
    "\n",
    "    if pattern == 'iid':\n",
    "        dict_users = iid(y_train, num_users)\n",
    "\n",
    "    elif pattern == 'noniid':\n",
    "        dict_users = noniid(y_train, num_users)\n",
    "    elif pattern > \"1\" and pattern <= \"9\":\n",
    "        #todo labelnoniid\n",
    "        # exit('Error: unfinsh')\n",
    "        dict_users=label(pattern,num_users,y_train)\n",
    "    elif pattern == \"iid-q\":\n",
    "        idxs = np.random.permutation(y_train.shape[0])\n",
    "        min_size = 0\n",
    "        while min_size < 10:\n",
    "            proportions = np.random.dirichlet(np.repeat(0.5, num_users))\n",
    "            proportions = proportions/proportions.sum()\n",
    "            min_size = np.min(proportions*len(idxs))\n",
    "        proportions = (np.cumsum(proportions)*len(idxs)).astype(int)[:-1]\n",
    "        batch= np.split(idxs,proportions)\n",
    "        dict_users = {i:batch[i] for i in range(num_users)}\n",
    "    else:\n",
    "        exit('Error: unrecognized pattern')\n",
    "        \n",
    "    \n",
    "    # all_idxs = [i for i in range(y_train.shape[0])]\n",
    "    # l=len(len(dataset_test))\n",
    "    # alp=0.8\n",
    "    # all_idxs_test = [i for i in range(l)]\n",
    "    # train_sampler = SubsetRandomSampler(all_idxs[:int(l*alp)])\n",
    "    # val_sampler = SubsetRandomSampler(all_idxs[int(l*alp):l])\n",
    "    # test_sampler = SubsetRandomSampler(all_idxs_test)\n",
    "\n",
    "    # print(dict_users)\n",
    "    train_dict = {}\n",
    "    val_dict={}\n",
    "    for i in range(num_users):\n",
    "        t,v = random_split(dict_users[i],lengths=[0.8,0.2],generator=torch.Generator().manual_seed(42))\n",
    "        train_dict[i] = t\n",
    "        val_dict[i] = v\n",
    "    return dict_users, train_dict, val_dict, dataset_train, dataset_test\n",
    "    # , dataset_test_part\n",
    "\n",
    "\n",
    "def iid(dataset, num_users):\n",
    "    \"\"\"\n",
    "    iid partition\n",
    "    \"\"\"\n",
    "    print(\"iid partion\")\n",
    "    idxs = np.random.permutation(dataset.shape[0])\n",
    "    batch = np.array_split(idxs,num_users)\n",
    "    dict_users = {i:batch[i] for i in range(num_users)}\n",
    "    return dict_users\n",
    "\n",
    "\n",
    "def noniid(dataset, num_users):\n",
    "    \"\"\"\n",
    "    noniid partition\n",
    "    dirichlet=0.5\n",
    "    \"\"\"\n",
    "    min_size = 0\n",
    "    min_require_size = 10\n",
    "    K = 10\n",
    "    dirichlet = 0.5\n",
    "    print(\"noniid partion\",dirichlet)\n",
    "    N = dataset.shape[0]\n",
    "    #np.random.seed(2020)\n",
    "    net_dataidx_map = {}\n",
    "\n",
    "    while min_size < min_require_size:\n",
    "        idx_batch = [[] for _ in range(num_users)]\n",
    "        for k in range(K):\n",
    "            idx_k = np.where(dataset== k)[0]\n",
    "            np.random.shuffle(idx_k)\n",
    "            proportions = np.random.dirichlet(np.repeat(dirichlet, num_users))\n",
    "            \n",
    "            ## Balance\n",
    "            proportions = np.array([p * (len(idx_j) < N / num_users) for p, idx_j in zip(proportions, idx_batch)])\n",
    "            proportions = proportions / proportions.sum()\n",
    "            proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]\n",
    "            idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))]\n",
    "            min_size = min([len(idx_j) for idx_j in idx_batch])\n",
    "    for j in range(num_users):\n",
    "        np.random.shuffle(idx_batch[j])\n",
    "        net_dataidx_map[j] = idx_batch[j]\n",
    "    return net_dataidx_map\n",
    "\n",
    "def label(partition:str,num_users,y_train):\n",
    "    num = eval(partition)\n",
    "    K = 10\n",
    "    if num == 10:\n",
    "        net_dataidx_map ={i:np.ndarray(0,dtype=np.int64) for i in range(num_users)}\n",
    "        for i in range(10):\n",
    "            idx_k = np.where(y_train==i)[0]\n",
    "            np.random.shuffle(idx_k)\n",
    "            split = np.array_split(idx_k,num_users)\n",
    "            for j in range(num_users):\n",
    "                net_dataidx_map[j]=np.append(net_dataidx_map[j],split[j])\n",
    "    else:\n",
    "        times=[0 for i in range(K)]\n",
    "        contain=[]\n",
    "        for i in range(num_users):\n",
    "            current=[i%K]\n",
    "            times[i%K]+=1\n",
    "            j=1\n",
    "            while (j<num):\n",
    "                ind=np.random.randint(0,K-1)\n",
    "                if (ind not in current):\n",
    "                    j=j+1\n",
    "                    current.append(ind)\n",
    "                    times[ind]+=1\n",
    "            contain.append(current)\n",
    "        net_dataidx_map ={i:np.ndarray(0,dtype=np.int64) for i in range(num_users)}\n",
    "        for i in range(K):\n",
    "            idx_k = np.where(y_train==i)[0]\n",
    "            np.random.shuffle(idx_k)\n",
    "            split = np.array_split(idx_k,times[i])\n",
    "            ids=0\n",
    "            for j in range(num_users):\n",
    "                if i in contain[j]:\n",
    "                    net_dataidx_map[j]=np.append(net_dataidx_map[j],split[ids])\n",
    "                    ids+=1\n",
    "    return net_dataidx_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#net\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, dim_in, dim_hidden, dim_out):\n",
    "        super(MLP, self).__init__()\n",
    "        self.layer_input = nn.Linear(dim_in, dim_hidden)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout()\n",
    "        self.layer_hidden = nn.Linear(dim_hidden, dim_out)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, x.shape[1] * x.shape[-2] * x.shape[-1])\n",
    "        x = self.layer_input(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.layer_hidden(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class CNN_Net(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(CNN_Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 64, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(64, 16, 7, 1)\n",
    "        self.fc1 = nn.Linear(4 * 4 * 16, 200)\n",
    "        self.fc2 = nn.Linear(200, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, 1, 32, 32)\n",
    "        x = torch.tanh(self.conv1(x))\n",
    "        x = F.max_pool2d(x, 2, 2)\n",
    "        x = torch.tanh(self.conv2(x))\n",
    "        x = F.max_pool2d(x, 2, 2)\n",
    "        x = x.view(-1, 4 * 4 * 16)\n",
    "        x = torch.tanh(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "\n",
    "\n",
    "class CNNMnist(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super(CNNMnist, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(args.num_channels, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, args.num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class CNNCifar(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super(CNNCifar, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, args.num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 5 * 5)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "class CNNfashion(nn.Module): # extend nn.Module class of nn\n",
    "    def __init__(self):\n",
    "        super().__init__() # super class constructor\n",
    "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=(5,5))\n",
    "        self.batchN1 = nn.BatchNorm2d(num_features=6)\n",
    "        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=(5,5))\n",
    "        self.fc1 = nn.Linear(in_features=12*4*4, out_features=120)\n",
    "        self.batchN2 = nn.BatchNorm1d(num_features=120)\n",
    "        self.fc2 = nn.Linear(in_features=120, out_features=60)\n",
    "        self.out = nn.Linear(in_features=60, out_features=10)\n",
    "        \n",
    "        \n",
    "        \n",
    "    def forward(self, x): # implements the forward method (flow of tensors)\n",
    "        \n",
    "        # hidden conv layer \n",
    "        x = self.conv1(x)\n",
    "        x = F.max_pool2d(input=x, kernel_size=2, stride=2)\n",
    "        x = F.relu(x)\n",
    "        x = self.batchN1(x)\n",
    "        \n",
    "        # hidden conv layer\n",
    "        x = self.conv2(x)\n",
    "        x = F.max_pool2d(input=x, kernel_size=2, stride=2)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        # flatten\n",
    "        x = x.reshape(-1, 12*4*4)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.batchN2(x)\n",
    "        x = self.fc2(x)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        # output\n",
    "        x = self.out(x)\n",
    "        \n",
    "        return x   \n",
    "\n",
    "import torch.nn as nn\n",
    "\n",
    "_cfg = {\n",
    "    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
    "    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
    "    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],\n",
    "    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],\n",
    "}\n",
    "\n",
    "\n",
    "def _make_layers(cfg):\n",
    "    layers = []\n",
    "    in_channels = 3\n",
    "    for layer_cfg in cfg:\n",
    "        if layer_cfg == 'M':\n",
    "            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))\n",
    "        else:\n",
    "            layers.append(nn.Conv2d(in_channels=in_channels,\n",
    "                                    out_channels=layer_cfg,\n",
    "                                    kernel_size=3,\n",
    "                                    stride=1,\n",
    "                                    padding=1,\n",
    "                                    bias=True))\n",
    "            layers.append(nn.BatchNorm2d(num_features=layer_cfg))\n",
    "            layers.append(nn.ReLU(inplace=True))\n",
    "            in_channels = layer_cfg\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "\n",
    "class _VGG(nn.Module):\n",
    "    \"\"\"\n",
    "    VGG module for 3x32x32 input, 10 classes\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, name):\n",
    "        super(_VGG, self).__init__()\n",
    "        cfg = _cfg[name]\n",
    "        self.layers = _make_layers(cfg)\n",
    "        flatten_features = 512\n",
    "        self.fc1 = nn.Linear(flatten_features, 10)\n",
    "        # self.fc2 = nn.Linear(4096, 4096)\n",
    "        # self.fc3 = nn.Linear(4096, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.layers(x)\n",
    "        y = y.view(y.size(0), -1)\n",
    "        y = self.fc1(y)\n",
    "        # y = self.fc2(y)\n",
    "        # y = self.fc3(y)\n",
    "        return y\n",
    "\n",
    "\n",
    "def VGG11():\n",
    "    return _VGG('VGG11')\n",
    "\n",
    "\n",
    "def VGG13():\n",
    "    return _VGG('VGG13')\n",
    "\n",
    "\n",
    "def VGG16():\n",
    "    return _VGG('VGG16')\n",
    "\n",
    "\n",
    "def VGG19():\n",
    "    return _VGG('VGG19')\n",
    "\n",
    "    def __init__(self):\n",
    "        super(VGG,self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3,64,3,padding=1)\n",
    "        self.conv2 = nn.Conv2d(64,64,3,padding=1)\n",
    "        self.pool1 = nn.MaxPool2d(2, 2)\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.relu1 = nn.ReLU()\n",
    "\n",
    "        self.conv3 = nn.Conv2d(64,128,3,padding=1)\n",
    "        self.conv4 = nn.Conv2d(128, 128, 3,padding=1)\n",
    "        self.pool2 = nn.MaxPool2d(2, 2, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(128)\n",
    "        self.relu2 = nn.ReLU()\n",
    "\n",
    "        self.conv5 = nn.Conv2d(128,128, 3,padding=1)\n",
    "        self.conv6 = nn.Conv2d(128, 128, 3,padding=1)\n",
    "        self.conv7 = nn.Conv2d(128, 128, 1,padding=1)\n",
    "        self.pool3 = nn.MaxPool2d(2, 2, padding=1)\n",
    "        self.bn3 = nn.BatchNorm2d(128)\n",
    "        self.relu3 = nn.ReLU()\n",
    "\n",
    "        self.conv8 = nn.Conv2d(128, 256, 3,padding=1)\n",
    "        self.conv9 = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.conv10 = nn.Conv2d(256, 256, 1, padding=1)\n",
    "        self.pool4 = nn.MaxPool2d(2, 2, padding=1)\n",
    "        self.bn4 = nn.BatchNorm2d(256)\n",
    "        self.relu4 = nn.ReLU()\n",
    "\n",
    "        self.conv11 = nn.Conv2d(256, 512, 3, padding=1)\n",
    "        self.conv12 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.conv13 = nn.Conv2d(512, 512, 1, padding=1)\n",
    "        self.pool5 = nn.MaxPool2d(2, 2, padding=1)\n",
    "        self.bn5 = nn.BatchNorm2d(512)\n",
    "        self.relu5 = nn.ReLU()\n",
    "\n",
    "        self.fc14 = nn.Linear(512*4*4,1024)\n",
    "        self.drop1 = nn.Dropout2d()\n",
    "        self.fc15 = nn.Linear(1024,1024)\n",
    "        self.drop2 = nn.Dropout2d()\n",
    "        self.fc16 = nn.Linear(1024,10)\n",
    "\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.pool1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu1(x)\n",
    "\n",
    "\n",
    "        x = self.conv3(x)\n",
    "        x = self.conv4(x)\n",
    "        x = self.pool2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu2(x)\n",
    "\n",
    "        x = self.conv5(x)\n",
    "        x = self.conv6(x)\n",
    "        x = self.conv7(x)\n",
    "        x = self.pool3(x)\n",
    "        x = self.bn3(x)\n",
    "        x = self.relu3(x)\n",
    "\n",
    "        x = self.conv8(x)\n",
    "        x = self.conv9(x)\n",
    "        x = self.conv10(x)\n",
    "        x = self.pool4(x)\n",
    "        x = self.bn4(x)\n",
    "        x = self.relu4(x)\n",
    "\n",
    "        x = self.conv11(x)\n",
    "        x = self.conv12(x)\n",
    "        x = self.conv13(x)\n",
    "        x = self.pool5(x)\n",
    "        x = self.bn5(x)\n",
    "        x = self.relu5(x)\n",
    "        # print(\" x shape \",x.size())\n",
    "        x = x.view(-1,512*4*4)\n",
    "        x = F.relu(self.fc14(x))\n",
    "        x = self.drop1(x)\n",
    "        x = F.relu(self.fc15(x))\n",
    "        x = self.drop2(x)\n",
    "        x = self.fc16(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, in_planes, planes, stride=1):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,\n",
    "                               stride=1, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != self.expansion*planes:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, self.expansion*planes,\n",
    "                          kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion*planes)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        out += self.shortcut(x)\n",
    "        out = F.relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, in_planes, planes, stride=1):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,\n",
    "                               stride=stride, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "        self.conv3 = nn.Conv2d(planes, self.expansion *\n",
    "                               planes, kernel_size=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(self.expansion*planes)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != self.expansion*planes:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, self.expansion*planes,\n",
    "                          kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion*planes)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = F.relu(self.bn2(self.conv2(out)))\n",
    "        out = self.bn3(self.conv3(out))\n",
    "        out += self.shortcut(x)\n",
    "        out = F.relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "# class ResNet(nn.Module):\n",
    "#     def __init__(self, block, num_blocks, num_classes=10):\n",
    "#         super(ResNet, self).__init__()\n",
    "#         self.in_planes = 64\n",
    "\n",
    "#         self.conv1 = nn.Conv2d(3, 64, kernel_size=3,\n",
    "#                                stride=1, padding=1, bias=False)\n",
    "#         self.bn1 = nn.BatchNorm2d(64)\n",
    "#         self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)\n",
    "#         self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)\n",
    "#         self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)\n",
    "#         self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)\n",
    "#         self.linear = nn.Linear(512*block.expansion, num_classes)\n",
    "\n",
    "#     def _make_layer(self, block, planes, num_blocks, stride):\n",
    "#         strides = [stride] + [1]*(num_blocks-1)\n",
    "#         layers = []\n",
    "#         for stride in strides:\n",
    "#             layers.append(block(self.in_planes, planes, stride))\n",
    "#             self.in_planes = planes * block.expansion\n",
    "#         return nn.Sequential(*layers)\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         out = F.relu(self.bn1(self.conv1(x)))\n",
    "#         out = self.layer1(out)\n",
    "#         out = self.layer2(out)\n",
    "#         out = self.layer3(out)\n",
    "#         out = self.layer4(out)\n",
    "#         out = F.avg_pool2d(out, 4)\n",
    "#         out = out.view(out.size(0), -1)\n",
    "#         out = self.linear(out)\n",
    "#         return out\n",
    "\n",
    "\n",
    "# def ResNet18():\n",
    "#     return ResNet(BasicBlock, [2, 2, 2, 2])\n",
    "\n",
    "\n",
    "# def ResNet34():\n",
    "#     return ResNet(BasicBlock, [3, 4, 6, 3])\n",
    "\n",
    "\n",
    "# def ResNet50():\n",
    "#     return ResNet(Bottleneck, [3, 4, 6, 3])\n",
    "\n",
    "\n",
    "# def ResNet101():\n",
    "#     return ResNet(Bottleneck, [3, 4, 23, 3])\n",
    "\n",
    "\n",
    "# def ResNet152():\n",
    "#     return ResNet(Bottleneck, [3, 8, 36, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#test\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "def test_img(net_g, datatest, args,sampler=None):\n",
    "    net_g.eval()\n",
    "    # testing\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    if sampler is not None :\n",
    "        data_loader = DataLoader(datatest, batch_size=args.bs,sampler=sampler)\n",
    "    else:\n",
    "        data_loader = DataLoader(datatest, batch_size=args.bs)\n",
    "    for idx, (data, target) in enumerate(data_loader):\n",
    "        if args.gpu != -1:\n",
    "            data, target = data.cuda(), target.cuda()\n",
    "        log_probs = net_g(data)\n",
    "        # sum up batch loss\n",
    "        test_loss += F.cross_entropy(log_probs, target, reduction='sum').item()\n",
    "        # get the index of the max log-probability\n",
    "        y_pred = log_probs.data.max(1, keepdim=True)[1]\n",
    "        correct += y_pred.eq(target.data.view_as(y_pred)).long().cpu().sum()\n",
    "\n",
    "    test_loss /= len(data_loader.dataset)\n",
    "    accuracy = correct.item() / 100\n",
    "\n",
    "    return accuracy, test_loss\n",
    "\n",
    "def test_dis(net,dataset,args):\n",
    "    net.eval()\n",
    "    data_loader = DataLoader(dataset,batch_size=args.bs)\n",
    "    y=[]\n",
    "    for data,target in data_loader:\n",
    "        data, target = data.cuda(), target.cuda()\n",
    "        predits = net(data)\n",
    "        y.append(predits.data.max(1)[1].cpu().numpy())\n",
    "    \n",
    "        \n",
    "    return Counter(np.concatenate(y))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#update\n",
    "import copy\n",
    "import math\n",
    "\n",
    "import torch\n",
    "from torch import nn, autograd\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from sklearn import metrics\n",
    "\n",
    "\n",
    "# class DatasetSplit(Dataset):\n",
    "#     def __init__(self, dataset, idxs):\n",
    "#         self.dataset = dataset\n",
    "#         self.idxs = list(idxs)\n",
    "\n",
    "#     def __len__(self):\n",
    "#         return len(self.idxs)\n",
    "\n",
    "#     def __getitem__(self, item):\n",
    "#         image, label = self.dataset[self.idxs[item]]\n",
    "#         return image, label\n",
    "\n",
    "\n",
    "class LocalUpdate(object):\n",
    "    def __init__(self, args, dataset=None, train=None):\n",
    "        self.args = args\n",
    "        self.loss_func = nn.CrossEntropyLoss()\n",
    "        self.selected_clients = []\n",
    "        self.len_train = len(train)\n",
    "        self.len_val = int(self.len_train*0.2)\n",
    "        self.ldr_train = DataLoader(dataset,batch_size=self.args.local_bs, drop_last=True,sampler=torch.utils.data.SubsetRandomSampler(train))\n",
    "        self.ldr_val = DataLoader(dataset,batch_size=self.args.local_bs, drop_last=True,sampler=torch.utils.data.SubsetRandomSampler(train[:self.len_val]))\n",
    "\n",
    "    def train(self, net):\n",
    "        net.train()\n",
    "        # train and update\n",
    "        optimizer = torch.optim.SGD(net.parameters(), lr=self.args.lr, momentum=self.args.momentum)\n",
    "\n",
    "        epoch_loss = []\n",
    "        for e in range(self.args.local_ep):\n",
    "            batch_loss = []\n",
    "            for batch_idx, (images, labels) in enumerate(self.ldr_train):\n",
    "                images, labels = images.to(self.args.device), labels.to(self.args.device)\n",
    "                net.zero_grad()  \n",
    "                log_probs = net(images)  \n",
    "                loss = self.loss_func(log_probs, labels)  \n",
    "                loss.backward() \n",
    "                optimizer.step()  \n",
    "                if self.args.verbose and batch_idx % 10 == 0:\n",
    "                    print('Update Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                        e, batch_idx * len(images), self.len_train,\n",
    "                              100. * batch_idx / self.len_train, loss.item()))\n",
    "                batch_loss.append(loss.item())\n",
    "            epoch_loss.append(sum(batch_loss) / len(batch_loss))\n",
    "        return net, net.state_dict(), sum(epoch_loss) / len(epoch_loss) \n",
    "    \n",
    "    def val(self,net, args):\n",
    "        net.eval()\n",
    "        # testing\n",
    "        test_loss = 0\n",
    "        correct = 0\n",
    "        # data_loader = DataLoader(datatest, batch_size=args.bs)\n",
    "        for idx, (data, target) in enumerate(self.ldr_val):\n",
    "            if args.gpu != -1:\n",
    "                data, target = data.cuda(), target.cuda()\n",
    "            log_probs = net(data)\n",
    "            \n",
    "            # sum up batch loss\n",
    "            test_loss += F.cross_entropy(log_probs, target, reduction='sum').item()\n",
    "            # get the index of the max log-probability\n",
    "            y_pred = log_probs.data.max(1, keepdim=True)[1]\n",
    "            correct += y_pred.eq(target.data.view_as(y_pred)).long().cpu().sum()\n",
    "        test_loss /= self.len_val\n",
    "        accuracy = correct / self.len_val\n",
    "\n",
    "        return accuracy, test_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#utility\n",
    "###flatten function\n",
    "def flatten(idx_dict):\n",
    "    return torch.concat([torch.flatten(idx_dict[key]) for key in idx_dict])\n",
    "\n",
    "def unflatten(flattened, normal_shape):\n",
    "    w_local = {}\n",
    "    for k in normal_shape:\n",
    "        n = len(normal_shape[k].view(-1))\n",
    "        w_local[k] = (flattened[:n].reshape(normal_shape[k].shape)).clone().detach()\n",
    "        flattened=flattened[n:]\n",
    "    return w_local\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#defences\n",
    "def multi_krum_defence(all_updates, n_attackers, multi_k=False):\n",
    "    candidates = []\n",
    "    candidate_indices = []\n",
    "    remaining_updates = all_updates\n",
    "    all_indices = np.arange(len(all_updates))\n",
    "\n",
    "    while len(remaining_updates) > 2 * n_attackers + 2:\n",
    "        torch.cuda.empty_cache()\n",
    "        distances = []\n",
    "        for update in remaining_updates:\n",
    "            distance = []\n",
    "            for update_ in remaining_updates:\n",
    "                distance.append(torch.norm((update - update_)) ** 2)\n",
    "            distance = torch.Tensor(distance).float()\n",
    "            distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)\n",
    "\n",
    "        distances = torch.sort(distances, dim=1)[0]\n",
    "        scores = torch.sum(distances[:, :len(remaining_updates) - 2 - n_attackers], dim=1)\n",
    "        indices = torch.argsort(scores)[:len(remaining_updates) - 2 - n_attackers]\n",
    "\n",
    "        candidate_indices.append(all_indices[indices[0].cpu().numpy()])\n",
    "        all_indices = np.delete(all_indices, indices[0].cpu().numpy())\n",
    "        candidates = remaining_updates[indices[0]][None, :] if not len(candidates) else torch.cat((candidates, remaining_updates[indices[0]][None, :]), 0)\n",
    "        remaining_updates = torch.cat((remaining_updates[:indices[0]], remaining_updates[indices[0] + 1:]), 0)\n",
    "        if not multi_k:\n",
    "            break\n",
    "    # print(len(remaining_updates))\n",
    "\n",
    "    aggregate = torch.mean(candidates, dim=0)\n",
    "\n",
    "    # return aggregate, np.array(candidate_indices)\n",
    "    return aggregate\n",
    "\n",
    "def tr_mean(all_updates, n_attackers):\n",
    "    sorted_updates = torch.sort(all_updates, 0)[0]\n",
    "    out = torch.mean(sorted_updates[n_attackers:-n_attackers], 0) if n_attackers else torch.mean(sorted_updates,0)\n",
    "    return out\n",
    "\n",
    "def bulyan(all_updates, n_attackers):\n",
    "    nusers = all_updates.shape[0]\n",
    "    bulyan_cluster = []\n",
    "    candidate_indices = []\n",
    "    remaining_updates = all_updates\n",
    "    all_indices = np.arange(len(all_updates))\n",
    "\n",
    "    while len(bulyan_cluster) < (nusers - 2 * n_attackers):\n",
    "        distances = []\n",
    "        for update in remaining_updates:\n",
    "            distance = torch.norm((remaining_updates - update), dim=1) ** 2\n",
    "            distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)\n",
    "\n",
    "        distances = torch.sort(distances, dim=1)[0]\n",
    "\n",
    "        scores = torch.sum(distances[:, :len(remaining_updates) - 2 - n_attackers], dim=1)\n",
    "        indices = torch.argsort(scores)[:len(remaining_updates) - 2 - n_attackers]\n",
    "\n",
    "        candidate_indices.append(all_indices[indices[0].cpu().numpy()])\n",
    "        all_indices = np.delete(all_indices, indices[0].cpu().numpy())\n",
    "        bulyan_cluster = remaining_updates[indices[0]][None, :] if not len(bulyan_cluster) else torch.cat((bulyan_cluster, remaining_updates[indices[0]][None, :]), 0)\n",
    "        remaining_updates = torch.cat((remaining_updates[:indices[0]], remaining_updates[indices[0] + 1:]), 0)\n",
    "\n",
    "    n, d = bulyan_cluster.shape\n",
    "    param_med = torch.median(bulyan_cluster, dim=0)[0]\n",
    "    sort_idx = torch.argsort(torch.abs(bulyan_cluster - param_med), dim=0)\n",
    "    sorted_params = bulyan_cluster[sort_idx, torch.arange(d)[None, :]]\n",
    "\n",
    "    # return torch.mean(sorted_params[:n - 2 * n_attackers], dim=0), np.array(candidate_indices)\n",
    "    return torch.mean(sorted_params[:n - 2 * n_attackers], dim=0)\n",
    "\n",
    "def dnc(updates,n_attackers):\n",
    "    d = len(updates[1])\n",
    "    num_iters = 1\n",
    "    sub_dim= 1000\n",
    "    fliter_frac=1.0\n",
    "    benign_ids = []\n",
    "    for i in range(num_iters):\n",
    "        indices = torch.randperm(d)[: sub_dim]\n",
    "        sub_updates = updates[:, indices]\n",
    "        mu = sub_updates.mean(dim=0)\n",
    "        centered_update = sub_updates - mu\n",
    "        v = torch.linalg.svd(centered_update, full_matrices=False)[2][0, :]\n",
    "        s = np.array(\n",
    "            [(torch.dot(update - mu, v) ** 2).item() for update in sub_updates]\n",
    "        )\n",
    "\n",
    "        good = s.argsort()[\n",
    "            : len(updates) - int(fliter_frac * n_attackers)\n",
    "        ]\n",
    "        benign_ids.extend(good)\n",
    "        print(benign_ids)\n",
    "    \n",
    "    #\n",
    "    benign_ids = list(set(benign_ids))\n",
    "    benign_updates = updates[benign_ids, :].mean(dim=0)\n",
    "    return benign_ids,benign_updates\n",
    "\n",
    "#\n",
    "def TDFL_cos(uw,t):\n",
    "    g,_=crh(uw,None)\n",
    "    cs=[]\n",
    "    for idx in range(len(uw)):\n",
    "        cs.append(torch.cosine_similarity(uw[idx],g,dim=0))\n",
    "    cs = torch.stack(cs)\n",
    "    print(cs)\n",
    "    return np.where(cs>=t)\n",
    "    \n",
    "###our \n",
    "##coming soon\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "##attack \n",
    "###fang attack\n",
    "def attack_median_and_trimmedmean(benign_update,m):\n",
    "    \n",
    "    # benign_update = \n",
    "    agg_grads = torch.mean(benign_update, 0)\n",
    "    deviation = torch.sign(agg_grads)\n",
    "    device = benign_update.device\n",
    "    b = 2\n",
    "    max_vector = torch.max(benign_update, 0)[0]\n",
    "    min_vector = torch.min(benign_update, 0)[0]\n",
    "\n",
    "    max_ = (max_vector > 0).type(torch.FloatTensor).to(device)\n",
    "    min_ = (min_vector < 0).type(torch.FloatTensor).to(device)\n",
    "\n",
    "    max_[max_ == 1] = b\n",
    "    max_[max_ == 0] = 1 / b\n",
    "    min_[min_ == 1] = b\n",
    "    min_[min_ == 0] = 1 / b\n",
    "\n",
    "    max_range = torch.cat(\n",
    "        (max_vector[:, None], (max_vector * max_)[:, None]), dim=1\n",
    "    )\n",
    "    min_range = torch.cat(\n",
    "        ((min_vector * min_)[:, None], min_vector[:, None]), dim=1\n",
    "    )\n",
    "\n",
    "    rand = (\n",
    "        torch.from_numpy(\n",
    "            np.random.uniform(0, 1, [len(deviation), m])\n",
    "        )\n",
    "        .type(torch.FloatTensor)\n",
    "        .to(benign_update.device)\n",
    "    )\n",
    "\n",
    "    max_rand = (\n",
    "        torch.stack([max_range[:, 0]] * rand.shape[1]).T\n",
    "        + rand * torch.stack([max_range[:, 1] - max_range[:, 0]] * rand.shape[1]).T\n",
    "    )\n",
    "    min_rand = (\n",
    "        torch.stack([min_range[:, 0]] * rand.shape[1]).T\n",
    "        + rand * torch.stack([min_range[:, 1] - min_range[:, 0]] * rand.shape[1]).T\n",
    "    )\n",
    "\n",
    "    mal_vec = (\n",
    "        torch.stack(\n",
    "            [(deviation < 0).type(torch.FloatTensor)] * max_rand.shape[1]\n",
    "        ).T.to(device)\n",
    "        * max_rand\n",
    "        + torch.stack(\n",
    "            [(deviation > 0).type(torch.FloatTensor)] * min_rand.shape[1]\n",
    "        ).T.to(device)\n",
    "        * min_rand\n",
    "    ).T\n",
    "    return mal_vec\n",
    "###\n",
    "def multi_krum(all_updates, n_attackers, multi_k=False):\n",
    "\n",
    "    candidates = []\n",
    "    candidate_indices = []\n",
    "    remaining_updates = all_updates\n",
    "    all_indices = np.arange(len(all_updates))\n",
    "\n",
    "    while len(remaining_updates) > 2 * n_attackers + 2:\n",
    "        torch.cuda.empty_cache()\n",
    "        distances = []\n",
    "        for update in remaining_updates:\n",
    "            distance = []\n",
    "            for update_ in remaining_updates:\n",
    "                distance.append(torch.norm((update - update_)) ** 2)\n",
    "            distance = torch.Tensor(distance).float()\n",
    "            distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)\n",
    "\n",
    "        distances = torch.sort(distances, dim=1)[0]\n",
    "        scores = torch.sum(distances[:, :len(remaining_updates) - 2 - n_attackers], dim=1)\n",
    "        indices = torch.argsort(scores)[:len(remaining_updates) - 2 - n_attackers]\n",
    "\n",
    "        candidate_indices.append(all_indices[indices[0].cpu().numpy()])\n",
    "        all_indices = np.delete(all_indices, indices[0].cpu().numpy())\n",
    "        candidates = remaining_updates[indices[0]][None, :] if not len(candidates) else torch.cat((candidates, remaining_updates[indices[0]][None, :]), 0)\n",
    "        remaining_updates = torch.cat((remaining_updates[:indices[0]], remaining_updates[indices[0] + 1:]), 0)\n",
    "        if not multi_k:\n",
    "            break\n",
    "    # print(len(remaining_updates))\n",
    "\n",
    "    aggregate = torch.mean(candidates, dim=0)\n",
    "\n",
    "    return aggregate, np.array(candidate_indices)\n",
    "\n",
    "def compute_lambda_fang(all_updates, model_re, n_attackers):\n",
    "\n",
    "    distances = []\n",
    "    n_benign, d = all_updates.shape\n",
    "    for update in all_updates:\n",
    "        distance = torch.norm((all_updates - update), dim=1)\n",
    "        distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)\n",
    "\n",
    "    distances[distances == 0] = 10000\n",
    "    distances = torch.sort(distances, dim=1)[0]\n",
    "    scores = torch.sum(distances[:, :n_benign - 2 - n_attackers], dim=1)\n",
    "    min_score = torch.min(scores)\n",
    "    term_1 = min_score / ((n_benign - n_attackers - 1) * torch.sqrt(torch.Tensor([d]))[0])\n",
    "    max_wre_dist = torch.max(torch.norm((all_updates - model_re), dim=1)) / (torch.sqrt(torch.Tensor([d]))[0])\n",
    "\n",
    "    return (term_1 + max_wre_dist)\n",
    "\n",
    "def get_malicious_updates_fang(all_updates, n_attackers):\n",
    "    model_re=torch.mean(all_updates,0)\n",
    "    deviation = torch.sign(model_re)\n",
    "    lamda = compute_lambda_fang(all_updates, model_re, n_attackers)\n",
    "    threshold = 1e-5\n",
    "\n",
    "    mal_updates = []    \n",
    "    while lamda > threshold:\n",
    "        mal_update = (- lamda * deviation)\n",
    "\n",
    "        mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "        mal_updates = torch.cat((mal_updates, all_updates), 0)\n",
    "\n",
    "        _, krum_candidate = multi_krum(mal_updates, n_attackers)\n",
    "        \n",
    "        if krum_candidate < n_attackers:\n",
    "            return mal_updates\n",
    "        \n",
    "        lamda *= 0.5\n",
    "\n",
    "    if not len(mal_updates):\n",
    "        print(lamda, threshold)\n",
    "        mal_update = (model_re - lamda * deviation)\n",
    "        \n",
    "        mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "        mal_updates = torch.cat((mal_updates, all_updates), 0)\n",
    "\n",
    "    return mal_updates[:n_attackers]\n",
    "###\n",
    "from scipy.stats import norm\n",
    "def LIE(para_list,m):\n",
    "    p=copy.deepcopy(para_list[:m])\n",
    "    n=para_list.shape[0]\n",
    "    \n",
    "    mu = torch.mean(para_list,0)\n",
    "    sigma = torch.std(para_list,0)\n",
    "    for i in range(m):\n",
    "        p[i]=mu-norm.ppf((n/2-1)/(n-m))*sigma\n",
    "    \n",
    "    return p \n",
    "###    \n",
    "def scaling_attack(para_list,m):\n",
    "    p=copy.deepcopy(para_list[:m])\n",
    "    factor= para_list.shape[1]\n",
    "    for i in range(m):\n",
    "        p[i]=para_list[i]*factor\n",
    "    \n",
    "    return p\n",
    "###\n",
    "def mean_attack(para_list,m):\n",
    "\n",
    "    return torch.stack([-para_list[i] for i in range(m)])\n",
    "\n",
    "def full_mean_attack(para_list,m):\n",
    "\n",
    "    p=copy.deepcopy(para_list[:m])\n",
    "\n",
    "    if m == para_list.shape[1]:\n",
    "        return mean_attack(para_list,m)\n",
    "    \n",
    "    all_sum=torch.sum(para_list,0)\n",
    "    m_para_sum= torch.sum(para_list[:m],0)\n",
    "    for i in range(m):\n",
    "        p[i]=((-all_sum- m_para_sum)/m)\n",
    "    return p\n",
    "#### attack for tailored\n",
    "def AGR_tailored_attack_on_krum(all_updates, n_attackers, dev_type='unit_vec'):\n",
    "    model_re = torch.mean(all_updates,0)\n",
    "    if dev_type == 'unit_vec':\n",
    "        deviation = model_re / torch.norm(model_re)\n",
    "    elif dev_type == 'sign':\n",
    "        deviation = torch.sign(model_re)\n",
    "    elif dev_type == 'std':\n",
    "        deviation = torch.std(all_updates, 0)\n",
    "\n",
    "    lamda = torch.Tensor([3.0])\n",
    "\n",
    "    threshold_diff = 1e-5\n",
    "    lamda_fail = lamda\n",
    "    lamda_succ = 0\n",
    "\n",
    "    while torch.abs(lamda_succ - lamda) > threshold_diff:\n",
    "        mal_update = (model_re - lamda * deviation)\n",
    "        mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "        mal_updates = torch.cat((mal_updates, all_updates), 0)\n",
    "\n",
    "        agg_grads, krum_candidate = multi_krum(mal_updates, n_attackers, multi_k=True)\n",
    "        if np.sum(krum_candidate < n_attackers) == n_attackers:\n",
    "            # print('successful lamda is ', lamda)\n",
    "            lamda_succ = lamda\n",
    "            lamda = lamda + lamda_fail / 2\n",
    "        else:\n",
    "            lamda = lamda - lamda_fail / 2\n",
    "\n",
    "        lamda_fail = lamda_fail / 2\n",
    "\n",
    "    mal_update = (model_re - lamda_succ * deviation)\n",
    "    mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "    return mal_updates\n",
    "\n",
    "def AGR_tailored_attack_on_median(all_updates, n_attackers, dev_type='unit_vec'):\n",
    "    model_re = torch.mean(all_updates,0)\n",
    "\n",
    "    if dev_type == 'unit_vec':\n",
    "        deviation = model_re / torch.norm(model_re)  # unit vector, dir opp to good dir\n",
    "    elif dev_type == 'sign':\n",
    "        deviation = torch.sign(model_re)\n",
    "    elif dev_type == 'std':\n",
    "        deviation = torch.std(all_updates, 0)\n",
    "\n",
    "    lamda = torch.Tensor([10.0])#compute_lambda_our(all_updates, model_re, n_attackers)\n",
    "\n",
    "    threshold_diff = 1e-5\n",
    "    prev_loss = -1\n",
    "    lamda_fail = lamda\n",
    "    lamda_succ = 0\n",
    "    iters = 0 \n",
    "    while torch.abs(lamda_succ - lamda) > threshold_diff:\n",
    "        mal_update = (model_re - lamda * deviation)\n",
    "        mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "        mal_updates = torch.cat((mal_updates, all_updates), 0)\n",
    "\n",
    "        agg_grads = torch.median(mal_updates, 0)[0]\n",
    "        \n",
    "        loss = torch.norm(agg_grads - model_re)\n",
    "        \n",
    "        if prev_loss < loss:\n",
    "            lamda_succ = lamda\n",
    "            lamda = lamda + lamda_fail / 2\n",
    "        else:\n",
    "            lamda = lamda - lamda_fail / 2\n",
    "\n",
    "        lamda_fail = lamda_fail / 2\n",
    "        prev_loss = loss\n",
    "        \n",
    "    mal_update = (model_re - lamda_succ * deviation)\n",
    "    mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "    return mal_updates\n",
    "\n",
    "def AGR_tailored_attack_on_trmean(all_updates, n_attackers, dev_type='unit_vec'):\n",
    "    model_re = torch.mean(all_updates,0)\n",
    "    if dev_type == 'unit_vec':\n",
    "        deviation = model_re / torch.norm(model_re)  # unit vector, dir opp to good dir\n",
    "    elif dev_type == 'sign':\n",
    "        deviation = torch.sign(model_re)\n",
    "    elif dev_type == 'std':\n",
    "        deviation = torch.std(all_updates, 0)\n",
    "\n",
    "    lamda = torch.Tensor([10.0]) #compute_lambda_our(all_updates, model_re, n_attackers)\n",
    "    # print(lamda)\n",
    "    threshold_diff = 1e-5\n",
    "    prev_loss = -1\n",
    "    lamda_fail = lamda\n",
    "    lamda_succ = 0\n",
    "    iters = 0 \n",
    "    while torch.abs(lamda_succ - lamda) > threshold_diff:\n",
    "        mal_update = (model_re - lamda * deviation)\n",
    "        mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "        mal_updates = torch.cat((mal_updates, all_updates), 0)\n",
    "\n",
    "        agg_grads = tr_mean(mal_updates, n_attackers)\n",
    "        \n",
    "        loss = torch.norm(agg_grads - model_re)\n",
    "        \n",
    "        if prev_loss < loss:\n",
    "            # print('successful lamda is ', lamda)\n",
    "            lamda_succ = lamda\n",
    "            lamda = lamda + lamda_fail / 2\n",
    "        else:\n",
    "            lamda = lamda - lamda_fail / 2\n",
    "\n",
    "        lamda_fail = lamda_fail / 2\n",
    "        prev_loss = loss\n",
    "        \n",
    "    mal_update = (model_re - lamda_succ * deviation)\n",
    "    mal_updates = torch.stack([mal_update] * n_attackers)\n",
    "\n",
    "    return mal_updates\n",
    "###attack for AGR unknow\n",
    "def min_max(all_updates,m,dev_type='unit_vec'):\n",
    "    model_re = torch.mean(all_updates,0)\n",
    "    if dev_type == 'unit_vec':\n",
    "        deviation = model_re / torch.norm(model_re)  # unit vector, dir opp to good dir\n",
    "    elif dev_type == 'sign':\n",
    "        deviation = torch.sign(model_re)\n",
    "    elif dev_type == 'std':\n",
    "        deviation = torch.std(all_updates, 0)\n",
    "\n",
    "    lamda = torch.Tensor([10]).float()\n",
    "    # print(lamda)\n",
    "    threshold_diff = 1e-5\n",
    "    lamda_fail = lamda\n",
    "    lamda_succ = 0\n",
    "    \n",
    "    distances = []\n",
    "    for update in all_updates:\n",
    "        distance = torch.norm((all_updates - update), dim=1) ** 2\n",
    "        distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)\n",
    "    \n",
    "    max_distance = torch.max(distances)\n",
    "    del distances\n",
    "\n",
    "    while torch.abs(lamda_succ - lamda) > threshold_diff:\n",
    "        mal_update = (model_re - lamda * deviation)\n",
    "        distance = torch.norm((all_updates - mal_update), dim=1) ** 2\n",
    "        max_d = torch.max(distance)\n",
    "        \n",
    "        if max_d <= max_distance:\n",
    "            # print('successful lamda is ', lamda)\n",
    "            lamda_succ = lamda\n",
    "            lamda = lamda + lamda_fail / 2\n",
    "        else:\n",
    "            lamda = lamda - lamda_fail / 2\n",
    "\n",
    "        lamda_fail = lamda_fail / 2\n",
    "\n",
    "    mal_update = (model_re - lamda_succ * deviation)\n",
    "    mal_updates = torch.stack([mal_update] * m)\n",
    "    return mal_updates\n",
    "\n",
    "def min_sum(all_updates,m,dev_type='unit_vec'):\n",
    "    model_re = torch.mean(all_updates,0)\n",
    "\n",
    "    if dev_type == 'unit_vec':\n",
    "        deviation = model_re / torch.norm(model_re)  # unit vector, dir opp to good dir\n",
    "    elif dev_type == 'sign':\n",
    "        deviation = torch.sign(model_re)\n",
    "    elif dev_type == 'std':\n",
    "        deviation = torch.std(all_updates, 0)\n",
    "    \n",
    "    lamda = torch.Tensor([10.0]).float()\n",
    "    # print(lamda)\n",
    "    threshold_diff = 1e-5\n",
    "    lamda_fail = lamda\n",
    "    lamda_succ = 0\n",
    "    \n",
    "    distances = []\n",
    "    for update in all_updates:\n",
    "        distance = torch.norm((all_updates - update), dim=1) ** 2\n",
    "        distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)\n",
    "    \n",
    "    scores = torch.sum(distances, dim=1)\n",
    "    min_score = torch.min(scores)\n",
    "    del distances\n",
    "\n",
    "    while torch.abs(lamda_succ - lamda) > threshold_diff:\n",
    "        mal_update = (model_re - lamda * deviation)\n",
    "        distance = torch.norm((all_updates - mal_update), dim=1) ** 2\n",
    "        score = torch.sum(distance)\n",
    "        \n",
    "        if score <= min_score:\n",
    "            # print('successful lamda is ', lamda)\n",
    "            lamda_succ = lamda\n",
    "            lamda = lamda + lamda_fail / 2\n",
    "        else:\n",
    "            lamda = lamda - lamda_fail / 2\n",
    "\n",
    "        lamda_fail = lamda_fail / 2\n",
    "\n",
    "    # print(lamda_succ)\n",
    "    mal_update = (model_re - lamda_succ * deviation)\n",
    "    mal_updates = torch.stack([mal_update] * m)\n",
    "\n",
    "    return mal_updates\n",
    "###     \n",
    "def KL(P,Q,mask=None):\n",
    "    eps = 0.0000001\n",
    "    d = (P+eps).log()-(Q+eps).log()\n",
    "    d = P*d\n",
    "    if mask !=None:\n",
    "        d = d*mask\n",
    "    return torch.sum(d)\n",
    "def CE(P,Q,mask=None):\n",
    "    return KL(P,Q,mask)+KL(1-P,1-Q,mask)\n",
    "\n",
    "def umap(output, target, data_batch, eps=0.0000001):\n",
    "    # start_idx = 0\n",
    "    # for param in test_net.parameters():\n",
    "    #     length = len(param.data.view(-1))\n",
    "    #     param.data = output[start_idx: start_idx + length].reshape(param.data.shape).cuda()\n",
    "    #     start_idx = start_idx + length\n",
    "    global update\n",
    "    test_net=copy.deepcopy(target)\n",
    "    test_net.load_state_dict(unflatten(output,update[0]))\n",
    "    # print(test_net(data_batch))\n",
    "    output_net = test_net(data_batch)\n",
    "    \n",
    "    target_net = target(data_batch)\n",
    "    # Normalize each vector by its norm\n",
    "    (n, d) = output_net.shape\n",
    "    output_net_norm = torch.sqrt(torch.sum(output_net ** 2, dim=1, keepdim=True))\n",
    "    output_net = output_net / (output_net_norm + eps)\n",
    "    output_net[output_net != output_net] = 0\n",
    "\n",
    "    target_net_norm = torch.sqrt(torch.sum(target_net ** 2, dim=1, keepdim=True))\n",
    "    target_net = target_net / (target_net_norm + eps)\n",
    "    target_net[target_net != target_net] = 0\n",
    "\n",
    "    # Calculate the cosine similarity\n",
    "    model_similarity = torch.mm(output_net, output_net.transpose(0, 1))\n",
    "    # model_similarity = model_similarity - torch.min(model_similarity,dim=1)[0].view(-1,1)\n",
    "    model_distance = 1-model_similarity #[0,2]\n",
    "    model_distance[range(n), range(n)] = 3\n",
    "    model_distance = model_distance - torch.min(model_distance, dim=1)[0].view(-1, 1)\n",
    "    model_distance[range(n), range(n)] = 0\n",
    "\n",
    "    model_similarity = 1-model_distance\n",
    "\n",
    "    target_similarity = torch.mm(target_net, target_net.transpose(0, 1))\n",
    "    target_distance = 1-target_similarity\n",
    "    target_distance[range(n), range(n)] = 3\n",
    "    target_distance = target_distance - torch.min(target_distance,dim=1)[0].view(-1,1)\n",
    "    target_distance[range(n), range(n)] = 0\n",
    "    target_similarity = 1 - target_distance\n",
    "\n",
    "\n",
    "    # Scale cosine similarity to 0..1\n",
    "    model_similarity = (model_similarity + 1.0) / 2.0\n",
    "    target_similarity = (target_similarity + 1.0) / 2.0\n",
    "\n",
    "    # Transform them into probabilities\n",
    "    model_similarity = model_similarity / torch.sum(model_similarity, dim=1, keepdim=True)\n",
    "    target_similarity = target_similarity / torch.sum(target_similarity, dim=1, keepdim=True)\n",
    "\n",
    "    # Calculate the KL-divergence\n",
    "    loss = CE(target_similarity,model_similarity)\n",
    "    # print(target_similarity,model_similarity)\n",
    "    # exit()\n",
    "    return loss\n",
    "\n",
    "def calculate_distance(w_glob,target):\n",
    "    distance=torch.norm(w_glob - target, p=2)\n",
    "    # print(\"now distance: \"+str(distance.item()))\n",
    "    return distance.item()\n",
    "\n",
    "def calculate_malicious(target_model,nusers,w_locals,choice,p,dataset_train):\n",
    "    train_loader = DataLoader(dataset_train, batch_size=args.local_bs,drop_last=True)\n",
    "    data_batch, _ = next(iter(train_loader))\n",
    "    data_batch=data_batch.cuda()\n",
    "    target_model_para=flatten(target_model.state_dict()).cpu()\n",
    "    min_dis = 1000000\n",
    "    while True:\n",
    "        w_locals_c = w_locals\n",
    "        # Three attack primitives\n",
    "        if choice == 1:\n",
    "            w0 = p * torch.randn_like(target_model_para)\n",
    "        elif choice == 2:\n",
    "            w0 = p * target_model_para\n",
    "        elif choice == 3:\n",
    "            # Note that here is a simulation in the full-knowledge setting\n",
    "            w0 = p * (target_model_para - torch.sum(w_locals, dim=0) / nusers)\n",
    "\n",
    "        # print(\"current p:\"+str(p))\n",
    "        while len(w_locals_c) < nusers:\n",
    "            w_locals_c = w0[None, :] if len(w_locals_c) == 0 else torch.cat((w_locals_c, w0[None, :]), 0)\n",
    "        w_glob = torch.mean(w_locals_c,dim=0)\n",
    "        \n",
    "        global loss_name\n",
    "        if loss_name=='umap':\n",
    "            now_dis = umap(w_glob, target_model, data_batch)\n",
    "        elif loss_name=='l2':\n",
    "            now_dis = calculate_distance(w_glob, target_model_para)\n",
    "        # print(now_dis)\n",
    "        # print('step 1:'+str(now_dis))\n",
    "        \n",
    "        #unflatten+load\n",
    "        # start_idx=0\n",
    "        # for param in test_net.parameters():\n",
    "        #     length=len(param.data.view(-1))\n",
    "        #     param.data = w_glob[start_idx: start_idx+length].reshape(param.data.shape).cuda()\n",
    "        #     start_idx = start_idx + length\n",
    "            \n",
    "        ##\n",
    "        decay=1.1\n",
    "        # print(p,now_dis,min_dis)\n",
    "        if now_dis<=min_dis:\n",
    "            min_dis=now_dis\n",
    "            p/=decay\n",
    "        else :\n",
    "            p*=decay\n",
    "            break\n",
    "\n",
    "    malicious_model= w0\n",
    "\n",
    "    return malicious_model\n",
    "\n",
    "\n",
    "\n",
    "def fang_adap(all_para,n):\n",
    "    l = 0.01\n",
    "    max=torch.max(all_para,dim=0)[0]\n",
    "    min=torch.min(all_para,dim=0)[0]\n",
    "    x,y = all_para.shape\n",
    "    m=np.zeros((n,y))\n",
    "    for i in range(n):\n",
    "        for y in range(y):\n",
    "            m[i][y]=np.random.uniform(min[y],max[y])\n",
    "    m=torch.from_numpy(m)\n",
    "    para = torch.concat([all_para,m])\n",
    "    v,_=crh(para)\n",
    "    V = torch.zeros_like(para[0])\n",
    "    for e in range(50):\n",
    "        v_hat ,w = crh(para)\n",
    "        # print(v_hat,w)\n",
    "        m_weight = w[x:]\n",
    "        for idx in range(n):\n",
    "            for j in range(y):\n",
    "                m[idx][j] += 2*l*(v_hat[j] - v[j])*m_weight[idx]/sum(w)\n",
    "        para = torch.concat([all_para,m])\n",
    "        if(abs(sum(v_hat-V))<1e-7 ):\n",
    "            print(\"converge\",e)\n",
    "            V=copy.deepcopy(v_hat)\n",
    "            break\n",
    "        V=copy.deepcopy(v_hat)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#init\n",
    "from torchvision.models import vgg11,resnet18,ResNet18_Weights\n",
    "args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')\n",
    "dict_users, train_dict, val_dict, dataset_train, dataset_test= Dataset_config(args.dataset, args.num_users, args.pattern)\n",
    "\n",
    "if args.model == 'cnn' and args.dataset == 'cifar':\n",
    "    net_glob = CNNCifar(args=args).to(args.device)\n",
    "elif args.model == 'cnn' and args.dataset == 'mnist':\n",
    "    net_glob = CNNMnist(args=args).to(args.device)\n",
    "elif args.model == 'mlp':\n",
    "    len_in = 1\n",
    "    img_size = dataset_train[0][0].shape\n",
    "    for x in img_size:\n",
    "        len_in *= x\n",
    "    net_glob = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)\n",
    "elif args.model == 'resnet' and args.dataset == 'cifar':\n",
    "    net_glob=resnet18(weights=ResNet18_Weights.DEFAULT)\n",
    "    num_features=net_glob.fc.in_features\n",
    "    net_glob.fc=nn.Linear(num_features,10)\n",
    "    model=net_glob.to(args.device)\n",
    "elif args.model == 'cnn' and args.dataset == 'fashion':\n",
    "    net_glob = CNNfashion().to(args.device)   \n",
    "elif args.model == 'vgg' and args.dataset == 'cifar':\n",
    "    net_glob = VGG11().to(args.device)\n",
    "else:\n",
    "    exit('Error: unrecognized model')\n",
    "print(net_glob)\n",
    "# writer = SummaryWriter(log_dir='./logs')\n",
    "# net_glob.requires_grad_()\n",
    "net_glob.train()\n",
    "\n",
    "# copy weights\n",
    "w_glob = net_glob.state_dict()\n",
    "\n",
    "# print(dataset_train.targets.numpy()[(dict_users[0])])\n",
    "# print(w_glob)\n",
    "avg_weights = torch.tensor([1/args.num_users for i in range(args.num_users)])\n",
    "loss_train = []\n",
    "val_acc_list, net_list = [], []\n",
    "#test\n",
    "test_acc=[]\n",
    "fedavg_acc=[]\n",
    "fedavg_loss=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##train\n",
    "local_acc=[]\n",
    "for e in range(5):\n",
    "    w_locals = []\n",
    "    locals_acc=[]\n",
    "    idxs_users = range(args.num_users)\n",
    "    pre_para = copy.deepcopy(net_glob.to(args.device).state_dict())\n",
    "    for idx in idxs_users:\n",
    "        args.local_ep = 10\n",
    "        net_glob.load_state_dict(pre_para)\n",
    "        local = LocalUpdate(args=args, dataset=dataset_train, train=dict_users[idx])\n",
    "        local_net, w, _ = local.train(net=copy.deepcopy(net_glob).to(args.device))\n",
    "        w_locals.append(copy.deepcopy(w))\n",
    "        # net_glob.load_state_dict(w)\n",
    "        local_test, _ = local.val(local_net,args)\n",
    "        # print(local_test)\n",
    "        locals_acc.append(local_test)\n",
    "    \n",
    "    net_glob.to(args.device)\n",
    "    global_para = copy.deepcopy(net_glob.to(args.device).state_dict())\n",
    "    #avg\n",
    "    global_para = unflatten(torch.mean(torch.stack([flatten(i).cpu() for i in w_locals]),dim=0),pre_para)\n",
    "\n",
    "    net_glob.load_state_dict(global_para)\n",
    "\n",
    "\n",
    "    acc_test, loss_test = test_img(net_glob, dataset_test, args)\n",
    "    test_acc.append(acc_test)\n",
    "    print(acc_test)\n",
    "    local_acc.append(locals_acc)\n",
    "    print(locals_acc)\n",
    "np.save('./local_acc.npy',local_acc)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#incentive\n",
    "\n",
    "e= [i.item() for i in local_acc[0]]\n",
    "# e= torch.stack([i for i in local_acc[0]])\n",
    "# np.sort(e)[1]\n",
    "def inc(l):\n",
    "    alpha=0.05\n",
    "    beta=160\n",
    "    reward=[]\n",
    "    acc = np.sort(l)[1]\n",
    "    s=sum(l)\n",
    "    for i in l:\n",
    "        if i <=acc:\n",
    "            reward.append(alpha*(acc-i)+beta*i/s)\n",
    "        else:\n",
    "            reward.append(beta*i/s)\n",
    "    return reward\n",
    "# print(e)\n",
    "for i in local_acc:\n",
    "    print(inc([j.item() for j in i]))\n",
    "# inc(e)\n",
    "# for epoch in range(5): \n",
    "#     local_acc(epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#only local train (for noniid & label flip)\n",
    "local_acc=[]\n",
    "loss_locals = []\n",
    "w_locals = []\n",
    "locals_acc=[]\n",
    "idxs_users = range(args.num_users)\n",
    "pre_para = copy.deepcopy(net_glob.to(args.device).state_dict())\n",
    "for idx in idxs_users:\n",
    "    args.local_ep = 50\n",
    "    net_glob.load_state_dict(pre_para)\n",
    "    local = LocalUpdate(args=args, dataset=dataset_train, train=train_dict[idx],val=val_dict[idx])\n",
    "    local_net,w, loss,size = local.train(net=copy.deepcopy(net_glob).to(args.device))\n",
    "    w_locals.append(copy.deepcopy(w))\n",
    "    loss_locals.append(copy.deepcopy(loss))\n",
    "    net_glob.load_state_dict(w)\n",
    "    val_acc,val_loss = local.val(local_net,args)\n",
    "    print(val_acc,val_loss)\n",
    "    \n",
    "dir = './mnist_noniid_mlp/0.9/'\n",
    "np.save(dir+'50e_val_dict',val_dict)\n",
    "np.save(dir+'50e_train_dict',train_dict)\n",
    "np.save(dir+'50e_w_locals',w_locals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for fmpa\n",
    "update=np.load('./result/cifar_cnn/noniid_0.5/50e_w_locals.npy',allow_pickle=True)[()]\n",
    "w_locals = torch.stack([flatten(i).cpu() for i in update])\n",
    "test_net = copy.deepcopy(net_glob)\n",
    "# dataset_train\n",
    "m=20\n",
    "# b=args.num_users-m\n",
    "# loss_name ='umap'\n",
    "# loss_name='l2'\n",
    "for loss_name in 'umap':\n",
    "    for choice in [1,2,3]:\n",
    "        for x in range(5):\n",
    "            print(loss_name,choice)\n",
    "            p=calculate_malicious(net_glob,args.num_users,w_locals[m:],choice,10,dataset_train)\n",
    "            poison=torch.vstack([p.reshape(1,-1)]*m)\n",
    "\n",
    "            uw = torch.vstack([poison,w_locals[m:]])\n",
    "\n",
    "            net_glob.load_state_dict(unflatten(torch.mean(w_locals[m:],dim=0),update[0]))\n",
    "            print(\"benign-avg\",test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "            net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "            print(\"poisoned-avg\",test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "            #dnc\n",
    "            benign_ids,final_w = dnc(uw,m)\n",
    "            print(\"Dnc\",benign_ids)\n",
    "            net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "            # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args))\n",
    "            print(\"Dnc\",test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "\n",
    "            g,distance= svtd(uw)\n",
    "            #tdfl\n",
    "            cs=[]\n",
    "            for idx in range(len(uw)):\n",
    "                cs.append(torch.cosine_similarity(uw[idx],g,dim=0))\n",
    "            cs = torch.stack(cs)\n",
    "            print(cs)\n",
    "            p= np.where(cs>=0.95)\n",
    "\n",
    "            print(\"tdfl\",p)\n",
    "            net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "            # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args))\n",
    "            print(\"tdfl\",test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "            # one_crh\n",
    "            distance = distance.numpy()\n",
    "            a=np.where(distance > np.mean(distance)-0.5)\n",
    "            if len(a[0])>(uw.shape[0]/2):\n",
    "                p1= a\n",
    "            else : p1= np.delete(np.arange(len(uw)),a)\n",
    "            net_glob.load_state_dict(unflatten(torch.mean(uw[p1],dim=0),update[0]))\n",
    "            print(\"oneTD\",test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "            #kmeanscrh\n",
    "            re = KMeans(2, random_state=0, n_init=\"auto\").fit(distance.reshape(-1,1)).labels_\n",
    "            print(re)\n",
    "            if len(np.where(re==0)[0])>len(np.where(re==1)[0]):\n",
    "                p2= np.where(re==0)[0]\n",
    "            else:\n",
    "                p2= np.where(re==1)[0]\n",
    "            net_glob.load_state_dict(unflatten(torch.mean(uw[p2],dim=0),update[0]))\n",
    "            print(\"kTD\",test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "#td\n",
    "# net_glob.load_state_dict(unflatten(torch.mean(one_crh(uw),dim=0),update[0]))\n",
    "# print(\"oneTD\",test_img(net_glob,dataset_test,args)[0])\n",
    "# net_glob.load_state_dict(unflatten(torch.mean(kmeans_crh(uw),dim=0),update[0]))\n",
    "# print(\"kTD\",test_img(net_glob,dataset_test,args)[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "ddir='./save/mnist_mlp_iid/'\n",
    "os.makedirs(ddir,exist_ok=True)\n",
    "\n",
    "update=np.load('/home/k3ats/self-fed/result/mnist_mlp/iid/50e_w_locals.npy',allow_pickle=True)[()]\n",
    "update_w = torch.stack([flatten(i).cpu() for i in update])\n",
    "m=20\n",
    "n_attacker = 20\n",
    "writer = pd.ExcelWriter(ddir+'/50epoch.xlsx')\n",
    "print('----------------table1----------------')\n",
    "table1_attack=[attack_median_and_trimmedmean,\n",
    "                get_malicious_updates_fang,\n",
    "                LIE,scaling_attack,\n",
    "                mean_attack,full_mean_attack,\n",
    "                min_max,min_sum]\n",
    "sheet1=[]    \n",
    "t1_record = []\n",
    "for attack in table1_attack:            \n",
    "    poison = attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])\n",
    "    record=[]\n",
    "    c_record = []\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    # print(attack.__name__,'mean',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    # median\n",
    "    net_glob.load_state_dict(unflatten(torch.median(uw,dim=0)[0],update[0]))\n",
    "    # print(attack.__name__,'median',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    #trim\n",
    "    final_w = tr_mean(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'trim',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #krum\n",
    "    if m ==24 :\n",
    "        final_w = multi_krum_defence(uw,20)\n",
    "    else:\n",
    "        final_w = multi_krum_defence(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'krum',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #bulyan\n",
    "    if(m==24):\n",
    "        final_w = bulyan(uw,20)\n",
    "    else:\n",
    "        final_w = bulyan(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'bulyan',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #dnc\n",
    "    benign_ids,final_w = dnc(uw,n_attacker)\n",
    "    c_record.append(benign_ids)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    \n",
    "    #TD\n",
    "    g,distance= crh(uw)\n",
    "    \n",
    "    #TDFL-cos\n",
    "    # p=TDFL_cos(uw,0.95)\n",
    "    cs=[]\n",
    "    for idx in range(len(uw)):\n",
    "        cs.append(torch.cosine_similarity(uw[idx],g,dim=0))\n",
    "    cs = torch.stack(cs)\n",
    "    print(cs)\n",
    "    p= np.where(cs>=0.95)\n",
    "    \n",
    "    c_record.append(p)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "    # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #one_crh\n",
    "    # p1=one_crh(uw)\n",
    "    distance = distance.numpy()\n",
    "    a=np.where(distance > np.mean(distance)-0.5)\n",
    "    if len(a[0])>(uw.shape[0]/2):\n",
    "        p1= a\n",
    "    else : p1= np.delete(np.arange(len(uw)),a)\n",
    "    \n",
    "    c_record.append(p1)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p1],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #kmeanscrh\n",
    "    # p2=kmeans_crh(uw)\n",
    "    \n",
    "    re = KMeans(2, random_state=0, n_init=\"auto\").fit(distance.reshape(-1,1)).labels_\n",
    "    print(re)\n",
    "    if len(np.where(re==0)[0])>len(np.where(re==1)[0]):\n",
    "        p2= np.where(re==0)[0]\n",
    "    else:\n",
    "        p2= np.where(re==1)[0]\n",
    "        \n",
    "    c_record.append(p2)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p2],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    sheet1.append(copy.deepcopy(record))\n",
    "    t1_record.append(copy.deepcopy(c_record))\n",
    "np.save(ddir+'/table1.npy',sheet1)\n",
    "np.save(ddir+'/t1_record.npy',t1_record)\n",
    "#save to excel\n",
    "f = pd.DataFrame(sheet1,\n",
    "                index=['attack_median_and_trimmedmean',\n",
    "                'get_malicious_updates_fang',\n",
    "                'LIE','scaling_attack',\n",
    "                'mean_attack','full_mean_attack','min_max','min_sum']\n",
    "                ,columns=['mean','median','tr_mean','krum','bulyan','dnc','TDFL_cos','one_crh','kmeans_crh']\n",
    "                ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table1\")\n",
    "\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.01\n",
      "0.05\n",
      "0.1\n",
      "0.3\n",
      "0.5\n",
      "0.7\n",
      "1\n",
      "          1      2      3      4      5     10     15     20     21     22   \n",
      "0.01  73.36  74.54  76.43  76.52  76.71  76.33  74.99  74.17  74.81  74.69  \\\n",
      "0.05  74.00  75.14  76.43  76.52  76.71  76.33  74.99  74.17  74.81  74.69   \n",
      "0.1   74.81  75.97  76.43  76.52  76.71  76.33  74.99  74.17  74.81  74.69   \n",
      "0.3   76.12  75.97  76.43  76.52  76.71  76.33  74.99  74.17  74.81  74.69   \n",
      "0.5   76.12  75.97  76.43  76.52  76.71  76.33  74.99  74.17  74.81  74.69   \n",
      "0.7   76.12  75.97  76.43  76.52  76.71  76.33  74.99  74.17  74.81  64.98   \n",
      "1     76.12  75.97  76.43  76.52  76.71  76.33  74.99  64.76  64.66  64.98   \n",
      "\n",
      "         23     24  \n",
      "0.01  73.85  73.74  \n",
      "0.05  73.85  73.74  \n",
      "0.1   73.85  73.74  \n",
      "0.3   73.85  64.58  \n",
      "0.5   65.21  64.58  \n",
      "0.7   65.21  64.58  \n",
      "1     65.21  64.58  \n"
     ]
    }
   ],
   "source": [
    "#threshold test\n",
    "update = np.load('/home/k3ats/self-fed/result/mnist_mlp/noniid_0.5/50e_w_locals.npy',allow_pickle=True)[()]\n",
    "ddir='./save/mnist_mlp_noniid_0.5/'\n",
    "os.makedirs(ddir,exist_ok=True)\n",
    "update_w = torch.stack([flatten(i).cpu() for i in update])\n",
    "writer = pd.ExcelWriter(ddir+'thres_record.xlsx')\n",
    "thres_record=[]\n",
    "for t in [0.01,0.05,0.1,0.3,0.5,0.7,1]:\n",
    "    record=[]\n",
    "    print(t)\n",
    "    for m in [1,2,3,4,5,10,15,20,21,22,23,24]:\n",
    "        poison = full_mean_attack(update_w,m)\n",
    "        uw = torch.vstack([poison,update_w[m:]])        \n",
    "        g,distance= crh(uw)\n",
    "        \n",
    "        # print(distance)\n",
    "        # break\n",
    "        distance = distance.numpy()\n",
    "        a=np.where(distance > np.mean(distance)-t)\n",
    "        if len(a[0])>(uw.shape[0]/2):\n",
    "            p1= a\n",
    "        else : p1= np.delete(np.arange(len(uw)),a)\n",
    "        # c_record.append(p1)\n",
    "        # print(p1)\n",
    "        net_glob.load_state_dict(unflatten(torch.mean(uw[p1],dim=0),update[0]))\n",
    "        # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args))\n",
    "        record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    thres_record.append(record)\n",
    "np.save(ddir+'./threshold_record.npy',thres_record)\n",
    "f = pd.DataFrame(thres_record,\n",
    "                 index=['0.01','0.05','0.1','0.3','0.5','0.7','1']\n",
    "                 ,columns=['1','2','3','4','5','10','15','20','21','22','23','24']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table1\")\n",
    "writer.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[[76.71, 77.37, 76.71, 76.71], [76.33, 76.83, 76.33, 76.33], [74.99, 71.35, 74.99, 74.99], [74.17, 44.56, 74.17, 74.17], [74.81, 41.23, 74.81, 74.81], [74.69, 35.25, 74.69, 74.69], [73.85, 30.65, 73.85, 73.85], [73.74, 25.31, 74.74, 73.74]]\n",
      "    withou attack  fedavg  One-TD   K-TD\n",
      "5           76.71   77.37   76.71  76.71\n",
      "10          76.33   76.83   76.33  76.33\n",
      "15          74.99   71.35   74.99  74.99\n",
      "20          74.17   44.56   74.17  74.17\n",
      "21          74.81   41.23   74.81  74.81\n",
      "22          74.69   35.25   74.69  74.69\n",
      "23          73.85   30.65   73.85  73.85\n",
      "24          73.74   25.31   74.74  73.74\n"
     ]
    }
   ],
   "source": [
    "#correct for one-TD\n",
    "update = np.load('./result/mnist_mlp/noniid_0.5/50e_w_locals.npy',allow_pickle=True)[()]\n",
    "ddir='./save/mnist_mlp_noniid_0.5/'\n",
    "os.makedirs(ddir,exist_ok=True)\n",
    "update_w = torch.stack([flatten(i).cpu() for i in update])\n",
    "writer = pd.ExcelWriter(ddir+'mean_ratio.xlsx')\n",
    "r_record=[]\n",
    "for m in [5,10,15,20,21,22,23,24]:\n",
    "    record=[]\n",
    "    poison = mean_attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])  \n",
    "\n",
    "    #without attack\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(update_w[m:],dim=0),update[0]))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "    g,distance= crh(uw)\n",
    "    # print(distance)\n",
    "    # break\n",
    "    distance = distance.numpy()\n",
    "    a=np.where(distance > np.mean(distance)-0.1)\n",
    "    if len(a[0])>(uw.shape[0]/2):\n",
    "        p1= a\n",
    "    else : p1= np.delete(np.arange(len(uw)),a)\n",
    "    # c_record.append(p1)\n",
    "    # print(p1)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p1],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "\n",
    "    re = KMeans(2, random_state=0, n_init=\"auto\").fit(distance.reshape(-1,1)).labels_\n",
    "    print(re)\n",
    "    if len(np.where(re==0)[0])>len(np.where(re==1)[0]):\n",
    "        p2= np.where(re==0)[0]\n",
    "    else:\n",
    "        p2= np.where(re==1)[0]\n",
    "        \n",
    "    # c_record.append(p2)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p2],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    r_record.append(record)\n",
    "np.save(ddir+'./mean_ratio.npy',r_record)\n",
    "print(r_record)\n",
    "f = pd.DataFrame(r_record,\n",
    "                 index=['5','10','15','20','21','22','23','24']\n",
    "                 ,columns=['withou attack','fedavg','One-TD','K-TD']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table1\")\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[[10.0, 52.92, 53.69, 49.66, 53.09, 53.75, 53.75], [10.0, 52.92, 53.26, 49.66, 52.48, 53.73, 53.73], [10.0, 52.23, 52.35, 49.66, 52.1, 53.67, 53.67], [10.0, 50.45, 51.32, 49.66, 10.0, 53.75, 53.75], [10.0, 39.63, 46.26, 48.55, 10.0, 53.7, 53.7]]\n",
      "    mean  median  trim-mean   krum  bulyan  One-TD   K-TD\n",
      "10  10.0   52.92      53.69  49.66   53.09   53.75  53.75\n",
      "20  10.0   52.92      53.26  49.66   52.48   53.73  53.73\n",
      "30  10.0   52.23      52.35  49.66   52.10   53.67  53.67\n",
      "40  10.0   50.45      51.32  49.66   10.00   53.75  53.75\n",
      "48  10.0   39.63      46.26  48.55   10.00   53.70  53.70\n"
     ]
    }
   ],
   "source": [
    "#for byzantine attack\n",
    "\n",
    "#mnist noniid\n",
    "update = np.load('./50c_resnet_cifar10_noniid0.5.npy',allow_pickle=True)[()]\n",
    "ddir='./save/cifar_resnet_noniid/'\n",
    "\n",
    "os.makedirs(ddir,exist_ok=True)\n",
    "update_w = torch.stack([flatten(i).cpu() for i in update])\n",
    "writer = pd.ExcelWriter(ddir+'m_attack.xlsx')\n",
    "r_record=[]\n",
    "for m in [5,10,15,20,24]:\n",
    "    n_attacker=m\n",
    "    record=[]\n",
    "    \n",
    "    poison = mean_attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])  \n",
    "\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    # print(attack.__name__,'mean',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    # median\n",
    "    net_glob.load_state_dict(unflatten(torch.median(uw,dim=0)[0],update[0]))\n",
    "    # print(attack.__name__,'median',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    #trim\n",
    "    final_w = tr_mean(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'trim',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #krum\n",
    "    if m==24: n_attacker=20\n",
    "    final_w = multi_krum_defence(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'krum',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    # #bulyan\n",
    "    if m==24: n_attacker=20\n",
    "    final_w = bulyan(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'bulyan',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #dnc\n",
    "    # benign_ids,final_w = dnc(uw,n_attacker)\n",
    "    # c_record.append(benign_ids)\n",
    "    # net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args))\n",
    "    # record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    \n",
    "    #TD\n",
    "    g,distance= crh(uw)\n",
    "    \n",
    "    #TDFL-cos\n",
    "    # p=TDFL_cos(uw,0.95)\n",
    "    # cs=[]\n",
    "    # for idx in range(len(uw)):\n",
    "    #     cs.append(torch.cosine_similarity(uw[idx],g,dim=0))\n",
    "    # cs = torch.stack(cs)\n",
    "    # print(cs)\n",
    "    # p= np.where(cs>=0.95)\n",
    "    \n",
    "    # c_record.append(p)\n",
    "    # net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "    # # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args))\n",
    "    # record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    # one_crh\n",
    "    # p1=one_crh(uw)\n",
    "    distance = distance.numpy()\n",
    "    a=np.where(distance > np.mean(distance)-0.1)\n",
    "    if len(a[0])>(uw.shape[0]/2):\n",
    "        p1= a\n",
    "    else : p1= np.delete(np.arange(len(uw)),a)\n",
    "    \n",
    "    c_record.append(p1)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p1],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #kmeanscrh\n",
    "    # p2=kmeans_crh(uw)\n",
    "    \n",
    "    re = KMeans(2, random_state=0, n_init=\"auto\").fit(distance.reshape(-1,1)).labels_\n",
    "    print(re)\n",
    "    if len(np.where(re==0)[0])>len(np.where(re==1)[0]):\n",
    "        p2= np.where(re==0)[0]\n",
    "    else:\n",
    "        p2= np.where(re==1)[0]\n",
    "        \n",
    "    c_record.append(p2)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p2],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    r_record.append(copy.deepcopy(record))\n",
    "np.save(ddir+'./m_attack.npy',r_record)\n",
    "print(r_record)\n",
    "f = pd.DataFrame(r_record,\n",
    "                 index=['10','20','30','40','48']\n",
    "                 ,columns=['mean','median','trim-mean','krum','bulyan','One-TD','K-TD']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table1\")\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exp\n",
    "import os\n",
    "import pandas as pd\n",
    "update = np.load('./result/mnist_mlp/iid/50e_w_locals.npy',allow_pickle=True)[()]\n",
    "ddir='./save/mnist_mlp_iid/'\n",
    "os.makedirs(ddir,exist_ok=True)\n",
    "\n",
    "update_w = torch.stack([flatten(i).cpu() for i in update])\n",
    "m=24\n",
    "n_attacker = 24\n",
    "# dir='./'\n",
    "writer = pd.ExcelWriter(ddir+'result.xlsx')\n",
    "print('----------------table1----------------')\n",
    "table1_attack=[AGR_tailored_attack_on_krum,AGR_tailored_attack_on_trmean,\n",
    "               attack_median_and_trimmedmean,get_malicious_updates_fang,LIE]\n",
    "sheet1=[] \n",
    "t1_record=[]   \n",
    "for attack in table1_attack:\n",
    "    poison = attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])\n",
    "    record=[]\n",
    "    c_record = []\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    # print(attack.__name__,'mean',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    # #trim\n",
    "    # final_w = tr_mean(uw,n_attacker)\n",
    "    # net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # # print(attack.__name__,'trim',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    # record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    # # median\n",
    "    # net_glob.load_state_dict(unflatten(torch.median(uw,dim=0)[0],update[0]))\n",
    "    # # print(attack.__name__,'median',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    # record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    #krum\n",
    "    # final_w = multi_krum_defence(uw,n_attacker)\n",
    "    # net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # # print(attack.__name__,'krum',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    # record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    # #bulyan\n",
    "    # final_w = bulyan(uw,n_attacker)\n",
    "    # net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # # print(attack.__name__,'bulyan',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    # record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #dnc\n",
    "    benign_ids,final_w = dnc(uw,n_attacker)\n",
    "    c_record.append(benign_ids)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    \n",
    "    #TD\n",
    "    g,distance= crh(uw)\n",
    "    \n",
    "    #TDFL-cos\n",
    "    # p=TDFL_cos(uw,0.95)\n",
    "    cs=[]\n",
    "    for idx in range(len(uw)):\n",
    "        cs.append(torch.cosine_similarity(uw[idx],g,dim=0))\n",
    "    cs = torch.stack(cs)\n",
    "    print(cs)\n",
    "    p= np.where(cs>=0.95)\n",
    "    \n",
    "    c_record.append(p)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "    # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    # one_crh\n",
    "    # p1=one_crh(uw)\n",
    "    distance = distance.numpy()\n",
    "    a=np.where(distance > np.mean(distance)-0.1)\n",
    "    if len(a[0])>(uw.shape[0]/2):\n",
    "        p1= a\n",
    "    else : p1= np.delete(np.arange(len(uw)),a)\n",
    "    \n",
    "    c_record.append(p1)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p1],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    #kmeanscrh\n",
    "    # p2=kmeans_crh(uw)\n",
    "    \n",
    "    re = KMeans(2, random_state=0, n_init=\"auto\").fit(distance.reshape(-1,1)).labels_\n",
    "    print(re)\n",
    "    if len(np.where(re==0)[0])>len(np.where(re==1)[0]):\n",
    "        p2= np.where(re==0)[0]\n",
    "    else:\n",
    "        p2= np.where(re==1)[0]\n",
    "        \n",
    "    c_record.append(p2)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p2],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args))\n",
    "    record.append(test_img(net_glob,dataset_test,args)[0])\n",
    "    \n",
    "    sheet1.append(copy.deepcopy(record))\n",
    "    t1_record.append(copy.deepcopy(c_record))\n",
    "np.save(ddir+'./table1.npy',sheet1)\n",
    "np.save(ddir+'./t1_record.npy',t1_record)\n",
    "#save to excel\n",
    "f = pd.DataFrame(sheet1,\n",
    "                 index=['AGR_tailored_attack_on_krum','AGR_tailored_attack_on_trmean','attack_median_and_trimmedmean',\n",
    "                        'get_malicious_updates_fang','LIE']\n",
    "                 ,columns=['mean','dnc','TDFL_cos','one_TD','K-TD']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table1\")\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "update = np.load('./50epoch_cifar_noniid0.5_50c.npy',allow_pickle=True)[()]\n",
    "\n",
    "\n",
    "update_w = torch.stack([flatten(i).cpu() for i in update])\n",
    "m=20\n",
    "n_attacker = 20\n",
    "dir='./cifar_noniid/'\n",
    "writer = pd.ExcelWriter(dir+'50epoch.xlsx')\n",
    "print('----------------table1----------------')\n",
    "table1_attack=[attack_median_and_trimmedmean,AGR_tailored_attack_on_trmean,LIE,scaling_attack,mean_attack,full_mean_attack,min_max,min_sum]\n",
    "sheet1=[]    \n",
    "for attack in table1_attack:\n",
    "    poison = attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])\n",
    "    record=[]\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    # print(attack.__name__,'mean',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    #trim\n",
    "    final_w = tr_mean(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'trim',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    #dnc\n",
    "    final_w = dnc(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    #TDFL-cos\n",
    "    p=TDFL_cos(uw,0.95)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "    # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #one_crh\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[one_crh(uw)],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #kmeanscrh\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[kmeans_crh(uw)],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    sheet1.append(copy.deepcopy(record))\n",
    "np.save(dir+'table1.npy',sheet1)\n",
    "#save to excel\n",
    "f = pd.DataFrame(sheet1,\n",
    "                 index=['attack_median_and_trimmedmean','AGR_tailored_attack_on_trmean','LIE','scaling_attack','mean_attack','full_mean_attack','min_max','min_sum']\n",
    "                 ,columns=['mean','tr_mean','dnc','TDFL_cos','one_crh','kmeans_crh']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table1\")\n",
    "\n",
    "#tabel2\n",
    "print('----------------table2----------------')\n",
    "sheet2=[]    \n",
    "table2_attack=[attack_median_and_trimmedmean,AGR_tailored_attack_on_median,LIE,scaling_attack,mean_attack,full_mean_attack,min_max,min_sum]\n",
    "for attack in table2_attack:\n",
    "    record=[]\n",
    "    poison = attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    # print(attack.__name__,'mean',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    # median\n",
    "    net_glob.load_state_dict(unflatten(torch.median(uw,dim=0)[0],update[0]))\n",
    "    # print(attack.__name__,'median',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #dnc\n",
    "    final_w = dnc(uw,n_attacker)    \n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #TDFL-cos\n",
    "    p=TDFL_cos(uw,0.95)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "    # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #one_crh\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[one_crh(uw)],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "        \n",
    "    #kmeanscrh\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[kmeans_crh(uw)],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    sheet2.append(record)\n",
    "np.save(dir+'table2.npy',sheet2)\n",
    "\n",
    "f = pd.DataFrame(sheet2,\n",
    "                 index=['attack_median_and_trimmedmean','AGR_tailored_attack_on_median','LIE','scaling_attack','mean_attack','full_mean_attack','min_max','min_sum']\n",
    "                 ,columns=['mean','median','dnc','TDFL_cos','one_crh','kmeans_crh']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table2\")\n",
    "\n",
    "#for table3   \n",
    "print('----------------table3----------------')\n",
    "table3_attack=[get_malicious_updates_fang,AGR_tailored_attack_on_krum,LIE,scaling_attack,mean_attack,full_mean_attack,min_max,min_sum]\n",
    "sheet3=[]     \n",
    "for attack in table3_attack:\n",
    "    record=[]\n",
    "    poison = attack(update_w,m)\n",
    "    uw = torch.vstack([poison,update_w[m:]])\n",
    "    #mean\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw,dim=0),update[0]))\n",
    "    # print(attack.__name__,'mean',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #krum\n",
    "    final_w = multi_krum_defence(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'krum',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #bulyan\n",
    "    final_w = bulyan(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'bulyan',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #dnc\n",
    "    final_w = dnc(uw,n_attacker)\n",
    "    net_glob.load_state_dict(unflatten(final_w,update[0]))\n",
    "    # print(attack.__name__,'dnc',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #TDFL-cos\n",
    "    p=TDFL_cos(uw,0.95)\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[p],dim=0),update[0]))\n",
    "    # print(attack.__name__,'TDFL_cos',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #one_crh\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[one_crh(uw)],dim=0),update[0]))\n",
    "    # print(attack.__name__,'one_crh',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    \n",
    "    #kmeanscrh\n",
    "    net_glob.load_state_dict(unflatten(torch.mean(uw[kmeans_crh(uw)],dim=0),update[0]))\n",
    "    # print(attack.__name__,'kmeans_crh',test_img(net_glob,dataset_test,args,test_sampler))\n",
    "    record.append(test_img(net_glob,dataset_test,args,test_sampler)[0])\n",
    "    sheet3.append(record)\n",
    "\n",
    "np.save(dir+'table3.npy',sheet3)\n",
    "f = pd.DataFrame(sheet3,\n",
    "                 index=['get_malicious_updates_fang','AGR_tailored_attack_on_krum','LIE','scaling_attack','mean_attack','full_mean_attack','min_max','min_sum']\n",
    "                 ,columns=['mean','krum','bulyan','dnc','TDFL_cos','one_crh','kmeans_crh']\n",
    "                 ).astype(np.float64)\n",
    "print(f)\n",
    "f.to_excel(writer,sheet_name=\"table3\")\n",
    "\n",
    "writer.close()\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}