FedMA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment if running on googlecolab \n",
    "# !pip install hickle\n",
    "# from google.colab import drive\n",
    "# drive.mount('/content/drive/')\n",
    "# %cd drive/MyDrive/PerCom2021-FL-master/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import copy\n",
    "from sklearn.preprocessing import normalize\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "import hickle as hkl \n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from lapsolver import solve_dense\n",
    "import csv\n",
    "import logging\n",
    "import numpy as np\n",
    "import os\n",
    "from sklearn.metrics import confusion_matrix, f1_score, precision_score, recall_score\n",
    "import argparse\n",
    "import json\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.utils.data as data\n",
    "from itertools import product\n",
    "import math\n",
    "import copy\n",
    "import time\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import confusion_matrix\n",
    "\n",
    "logging.basicConfig()\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "\n",
    "args_logdir = \"logs/cifar10\"\n",
    "#args_dataset = \"cifar10\"\n",
    "args_datadir = \"./data/cifar10\"\n",
    "args_init_seed = 0\n",
    "args_net_config = [3072, 100, 8]\n",
    "#args_partition = \"hetero-dir\"\n",
    "args_partition = \"homo\"\n",
    "args_experiment = [\"u-ensemble\", \"pdm\"]\n",
    "args_trials = 1\n",
    "#args_lr = 0.01\n",
    "args_epochs = 10\n",
    "args_reg = 1e-5\n",
    "args_alpha = 0.5\n",
    "args_iter_epochs = None\n",
    "\n",
    "args_pdm_sig = 1.0\n",
    "args_pdm_sig0 = 1.0\n",
    "args_pdm_gamma = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# which GPU to use\n",
    "# \"-1,0,1\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "# DNN,CNN\n",
    "modelType = \"CNN\"\n",
    "\n",
    "#prior informatin for model type \n",
    "layerType = [0,1,1]\n",
    "\n",
    "# algorithm .to('cpu)= \"FEDAVG,FEDPER\"\n",
    "algorithm = \"FEDMA\"\n",
    "\n",
    "# REALWORLD_CLIENT\n",
    "dataSetName = 'REALWORLD_CLIENT'\n",
    "\n",
    "#BALANCED, UNBALANCED\n",
    "dataConfig = \"BALANCED\"\n",
    "\n",
    "#ADAM, SGD\n",
    "optimizer = \"SGD\"\n",
    "\n",
    "#0, 1\n",
    "ClientAllTest = True\n",
    "\n",
    "# Save the client models a .h5 file\n",
    "savedClientModel = 0\n",
    "\n",
    "# Show training verbose\n",
    "showTrainVerbose = 0\n",
    "\n",
    "# input window size \n",
    "segment_size = 128\n",
    "\n",
    "# input channel count\n",
    "num_input_channels = 6\n",
    "\n",
    "# client learning rate\n",
    "learningRate = 0.01\n",
    "\n",
    "# model drop out rate\n",
    "dropout_rate = 0.5\n",
    "\n",
    "# local epoch\n",
    "localEpoch = 5\n",
    "\n",
    "# communication round\n",
    "communicationRound = 200\n",
    "\n",
    "# CNN kernal size\n",
    "kernelSize = 16\n",
    "\n",
    "# Neuron distance measurement \n",
    "euclid = True\n",
    "\n",
    "# Seed for data partioning and PyTorch training\n",
    "randomSeed = 1\n",
    "\n",
    "# FedMA parameter\n",
    "\n",
    "iteration = 5\n",
    "\n",
    "gammaValue = float(iteration + 2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if(dataSetName == 'UCI'):\n",
    "    ACTIVITY_LABEL = ['WALKING', 'WALKING_UPSTAIRS','WALKING_DOWNSTAIRS', 'SITTING', 'STANDING', 'LAYING']\n",
    "else:\n",
    "    ACTIVITY_LABEL = ['climbingdown', 'climbingup', 'jumping','lying', 'running', 'sitting', 'standing', 'walking']\n",
    "    \n",
    "activityCount = len(ACTIVITY_LABEL)\n",
    "\n",
    "if(modelType == \"DNN\"):\n",
    "    architectureType = str(algorithm)+'_'+str(learningRate)+'LR_'+str(localEpoch)+'LE_'+str(communicationRound)+'CR_400D_100D_'+str(dataSetName)+\"_IT\"+str(iteration)+\"_NOTSAME\"\n",
    "else: \n",
    "    architectureType = str(algorithm)+'_'+str(learningRate)+'LR_'+str(localEpoch)+'LE_'+str(communicationRound)+'CR_196-16C_4M_1024D_'+str(dataSetName)+\"_IT\"+str(iteration)+\"__NOTSAME\"\n",
    "mainDir = ''\n",
    "filepath = mainDir + 'savedModels/'+architectureType+'/'+dataSetName+'/'\n",
    "os.makedirs(filepath, exist_ok=True)\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'\n",
    "\n",
    "if(dataSetName=='UCI'):\n",
    "    clientCount = 5\n",
    "else:\n",
    "    clientCount = 15\n",
    "    \n",
    "np.random.seed(randomSeed)\n",
    "torch.manual_seed(randomSeed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "parser = argparse.ArgumentParser(description='Probabilistic Federated CNN Matching')\n",
    "\n",
    "parser.add_argument('--model', type=str, default='lenet', metavar='N',\n",
    "                    help='neural network used in training')\n",
    "parser.add_argument('--dataset', type=str, default='cifar10', metavar='N',\n",
    "                    help='dataset used for training')\n",
    "parser.add_argument('--partition', type=str, default='homo', metavar='N',\n",
    "                    help='how to partition the dataset on local workers')\n",
    "parser.add_argument('--batch-size', type=int, default=128, metavar='N',\n",
    "                    help='input batch size for training (default: 64)')\n",
    "parser.add_argument('--lr', type=float, default=0.1, metavar='LR',\n",
    "                    help='learning rate (default: 0.01)')\n",
    "parser.add_argument('--retrain_lr', type=float, default=0.1, metavar='RLR',\n",
    "                    help='learning rate using in specific for local network retrain (default: 0.01)')\n",
    "parser.add_argument('--fine_tune_lr', type=float, default=0.1, metavar='FLR',\n",
    "                    help='learning rate using in specific for fine tuning the softmax layer on the data center (default: 0.01)')\n",
    "parser.add_argument('--epochs', type=int, default=5, metavar='EP',\n",
    "                    help='how many epochs will be trained in a training process')\n",
    "parser.add_argument('--retrain_epochs', type=int, default=10, metavar='REP',\n",
    "                    help='how many epochs will be trained in during the locally retraining process')\n",
    "parser.add_argument('--fine_tune_epochs', type=int, default=10, metavar='FEP',\n",
    "                    help='how many epochs will be trained in during the fine tuning process')\n",
    "parser.add_argument('--partition_step_size', type=int, default=6, metavar='PSS',\n",
    "                    help='how many groups of partitions we will have')\n",
    "parser.add_argument('--local_points', type=int, default=5000, metavar='LP',\n",
    "                    help='the approximate fixed number of data points we will have on each local worker')\n",
    "parser.add_argument('--partition_step', type=int, default=0, metavar='PS',\n",
    "                    help='how many sub groups we are going to use for a particular training process')\n",
    "parser.add_argument('--n_nets', type=int, default=2, metavar='NN',\n",
    "                    help='number of workers in a distributed cluster')\n",
    "parser.add_argument('--oneshot_matching', type=bool, default=False, metavar='OM',\n",
    "                    help='if the code is going to conduct one shot matching')\n",
    "parser.add_argument('--retrain', type=bool, default=False,\n",
    "                    help='whether to retrain the model or load model locally')\n",
    "parser.add_argument('--rematching', type=bool, default=False,\n",
    "                    help='whether to recalculating the matching process (this is for speeding up the debugging process)')\n",
    "parser.add_argument('--comm_type', type=str, default='layerwise',\n",
    "                    help='which type of communication strategy is going to be used: layerwise/blockwise')\n",
    "parser.add_argument('--comm_round', type=int, default=10,\n",
    "                    help='how many round of communications we shoud use')\n",
    "\n",
    "\n",
    "temp = ['--model=simple-cnn',\n",
    " '--dataset=cifar10',\n",
    " '--lr='+str(learningRate),\n",
    " '--retrain_lr='+str(learningRate),\n",
    " '--batch-size=64',\n",
    " '--epochs='+str(localEpoch),\n",
    " '--retrain_epochs='+str(localEpoch),\n",
    " '--n_nets='+str(clientCount),\n",
    " '--partition=hetero-dir',\n",
    " '--comm_type=fedma',\n",
    " '--comm_round='+str(communicationRound),\n",
    " '--retrain=True',\n",
    " '--rematching=True']\n",
    "args = parser.parse_args(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dims, output_dim=activityCount):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv_layer = nn.Sequential(\n",
    "            nn.Conv1d(num_input_channels, 196, kernel_size=kernelSize, bias=True),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(4),\n",
    "        )\n",
    "        self.fc_layer = nn.Sequential(\n",
    "            nn.Linear(5488, 1024),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Dropout(p=dropout_rate),\n",
    "            nn.Linear(1024, output_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layer(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc_layer(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class SimpleCNNContainer(nn.Module):\n",
    "    def __init__(self, input_channel, num_filters, kernel_size, input_dim, hidden_dims, output_dim=activityCount):\n",
    "        super(SimpleCNNContainer, self).__init__()\n",
    "        self.conv_layer = nn.Sequential(\n",
    "            nn.Conv1d(input_channel, num_filters[0], kernel_size),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(4),\n",
    "        )\n",
    "        self.neurons = self.linear_input_neurons()\n",
    "#         logger.info(\"neurons {}\".format(self.neurons))\n",
    "        self.fc_layer = nn.Sequential(\n",
    "            nn.Linear(self.neurons, hidden_dims[0]),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Dropout(p=dropout_rate),\n",
    "            nn.Linear(hidden_dims[0],output_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layer(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "#         logger.info(\"x.size(0) {}\".format(x.size(0)))\n",
    "        x = self.fc_layer(x)\n",
    "        return x\n",
    "\n",
    "    def forward_conv(self, x):\n",
    "        x = self.conv_layer(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        return x\n",
    "\n",
    "    def size_after_relu(self, x):\n",
    "        x = self.conv_layer(x)\n",
    "        return x.size()\n",
    "\n",
    "    def linear_input_neurons(self):\n",
    "        size = self.size_after_relu(torch.rand(1, 6, 128))  # image size: 64x32\n",
    "        m = 1\n",
    "        for i in size:\n",
    "            m *= i\n",
    "        return int(m)\n",
    "\n",
    "    \n",
    "class SimpleCNNContainerConvBlocks(nn.Module):\n",
    "    def __init__(self, input_channel, num_filters, kernel_size, output_dim=activityCount):\n",
    "        super(SimpleCNNContainerConvBlocks, self).__init__()\n",
    "        \n",
    "        self.conv_layer = nn.Sequential(\n",
    "            nn.Conv1d(input_channel, num_filters[0], kernel_size),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(4),\n",
    "        )\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layer(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clientDataTrain = []\n",
    "clientLabelTrain = []\n",
    "clientDataTest = []\n",
    "clientLabelTest = []\n",
    "\n",
    "centralTrainData = []\n",
    "centralTrainLabel = []\n",
    "\n",
    "centralTestData = []\n",
    "centralTestLabel = []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if(dataSetName == \"UCI\"):\n",
    "    def load_file(filepath):\n",
    "        dataframe = pd.read_csv(filepath, header=None)\n",
    "        return dataframe.values\n",
    "\n",
    "\n",
    "    def load_group(filenames, prefix=''):\n",
    "        loaded = list()\n",
    "        for name in filenames:\n",
    "            data = load_file(prefix + name)\n",
    "            loaded.append(data)\n",
    "        loaded = np.dstack(loaded)\n",
    "        return loaded\n",
    "\n",
    "\n",
    "    def load_dataset(group, prefix=''):\n",
    "        filepath = mainDir + 'datasetStandardized/'+prefix + '/' + group + '/'\n",
    "        filenames = list()\n",
    "        filenames += ['AccX'+prefix+'.csv', 'AccY' +\n",
    "                      prefix+'.csv', 'AccZ'+prefix+'.csv']\n",
    "        filenames += ['GyroX'+prefix+'.csv', 'GyroY' +\n",
    "                      prefix+'.csv', 'GyroZ'+prefix+'.csv']\n",
    "        X = load_group(filenames, filepath)\n",
    "        y = load_file(mainDir + 'datasetStandardized/'+prefix +\n",
    "                      '/' + group + '/Label'+prefix+'.csv')\n",
    "        return X, y\n",
    "    trainData, trainLabel = load_dataset('train', dataSetName)\n",
    "    evalData, evalLabel = load_dataset('eval', dataSetName)\n",
    "    allData = np.float32(np.vstack((trainData, evalData)))\n",
    "    allLabel = np.vstack((trainLabel, evalLabel))\n",
    "\n",
    "    # split data into 80 - 20 \n",
    "    skf = StratifiedKFold(n_splits=5,shuffle = True)\n",
    "    skf.get_n_splits(allData, allLabel)\n",
    "    partitionedData = list()\n",
    "    partitionedLabel = list()\n",
    "    for train_index, test_index in skf.split(allData, allLabel):\n",
    "        partitionedData.append(allData[test_index])\n",
    "        partitionedLabel.append(allLabel[test_index])\n",
    "\n",
    "    centralTrainData = np.vstack((partitionedData[:4]))\n",
    "    centralTrainLabel = np.vstack((partitionedLabel[:4]))\n",
    "    centralTestData = partitionedData[4]\n",
    "    centralTestLabel = partitionedLabel[4]\n",
    "\n",
    "    trainData = list()\n",
    "    trainLabel = list()\n",
    "    testData = list()\n",
    "    testLabel = list()\n",
    "\n",
    "    if(dataConfig == \"BALANCED\"):\n",
    "        skf = StratifiedKFold(n_splits=clientCount)\n",
    "        skf.get_n_splits(centralTrainData, centralTrainLabel)\n",
    "        for train_index, test_index in skf.split(centralTrainData, centralTrainLabel):\n",
    "            trainData.append(np.asarray(centralTrainData[test_index]).reshape(-1,6,128))\n",
    "            trainLabel.append(centralTrainLabel[test_index].ravel().astype(int))\n",
    "    else:\n",
    "    # unbalanced\n",
    "        kf = KFold(n_splits=clientCount, shuffle=True)\n",
    "        kf.get_n_splits(centralTrainData)\n",
    "        for train_index, test_index in kf.split(centralTrainData):\n",
    "            trainData.append(np.asarray(centralTrainData[test_index]).reshape(-1,6,128))\n",
    "            trainLabel.append(centralTrainLabel[test_index].ravel().astype(int))\n",
    "\n",
    "    #slittestSetInto5\n",
    "    skf.get_n_splits(centralTestData, centralTestLabel)\n",
    "    for train_index, test_index in skf.split(centralTestData, centralTestLabel):\n",
    "        testData.append(np.asarray(centralTestData[test_index]).reshape(-1,6,128)[:411])\n",
    "        testLabel.append(centralTestLabel[test_index].ravel().astype(int)[:411])\n",
    "\n",
    "    clientDataTrain = np.vstack(trainData) \n",
    "    clientLabelTrain = trainLabel\n",
    "    clientDataTest = np.vstack(testData)\n",
    "    clientLabelTest = testLabel\n",
    "    \n",
    "    centralTrainData = np.float32(np.vstack(clientDataTrain))\n",
    "    centralTrainLabel = np.vstack(clientLabelTrain).ravel()\n",
    "    centralTestData = np.float32(np.vstack(clientDataTest))\n",
    "    centralTestLabel = np.vstack(clientLabelTest).ravel()\n",
    "    \n",
    "else:\n",
    "    clientData = []\n",
    "    clientLabel = []\n",
    "\n",
    "    dataSetName = 'REALWORLD_CLIENT'\n",
    "    for i in range(0,15):\n",
    "        accX = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/AccX'+dataSetName+'.hkl')\n",
    "        accY = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/AccY'+dataSetName+'.hkl')\n",
    "        accZ = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/AccZ'+dataSetName+'.hkl')\n",
    "        gyroX = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/GyroX'+dataSetName+'.hkl')\n",
    "        gyroY = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/GyroY'+dataSetName+'.hkl')\n",
    "        gyroZ = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/GyroZ'+dataSetName+'.hkl')\n",
    "        label = hkl.load('datasetStandardized/'+dataSetName+'/'+str(i)+'/Label'+dataSetName+'.hkl')\n",
    "        clientData.append(np.dstack((accX,accY,accZ,gyroX,gyroY,gyroZ)))\n",
    "        clientLabel.append(label)\n",
    "    \n",
    "    if(dataConfig == \"BALANCED\"):\n",
    "        for i in range (0,int(args.n_nets)):\n",
    "            skf = StratifiedKFold(n_splits=5, shuffle=True)\n",
    "            skf.get_n_splits(clientData[i], clientLabel[i])\n",
    "            partitionedData = list()\n",
    "            partitionedLabel = list()    \n",
    "            for train_index, test_index in skf.split(clientData[i], clientLabel[i]):\n",
    "                partitionedData.append(clientData[i][test_index])\n",
    "                partitionedLabel.append(clientLabel[i][test_index])\n",
    "            clientDataTrain.append(np.float32(np.vstack(partitionedData[:4])).reshape(-1,6,128))\n",
    "            clientLabelTrain.append((np.hstack((partitionedLabel[:4]))))\n",
    "            clientDataTest.append(np.float32(partitionedData[4]).reshape(-1,6,128))\n",
    "            clientLabelTest.append((partitionedLabel[4]))\n",
    "            \n",
    "    centralTrainData = np.float32(np.vstack(clientDataTrain))\n",
    "    centralTrainLabel = (np.hstack(clientLabelTrain))\n",
    "\n",
    "    centralTestData = np.float32(np.vstack((clientDataTest)))\n",
    "    centralTestLabel = (np.hstack(clientLabelTest))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traindata_cls_counts = {}\n",
    "net_dataidx_map = {}\n",
    "\n",
    "startingIndex = 0 \n",
    "endingIndex = 0\n",
    "y_train = []\n",
    "for i in range(0,args.n_nets):\n",
    "    traindata_cls_counts[i] = pd.Series(clientLabelTrain[i]).value_counts().sort_index().to_dict()\n",
    "    startingIndex = endingIndex\n",
    "    endingIndex = endingIndex + clientDataTrain[i].shape[0]\n",
    "    y_train.append(clientLabelTrain[i])\n",
    "    net_dataidx_map[i] = list(range(startingIndex, endingIndex))\n",
    "    \n",
    "y_train = np.hstack((y_train))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloader(dataset, datadir, train_bs, test_bs, dataidxs=None,clientNum=100):\n",
    "    if dataset in ('mnist', 'cifar10'):\n",
    "        dl_obj = CIFAR10_truncated\n",
    "        if(clientNum == 100):\n",
    "            train_ds = dl_obj(datadir,centralTrainData,centralTrainLabel, \n",
    "                          train=True,download=False)\n",
    "            test_ds = dl_obj(datadir,centralTestData,centralTestLabel,train=False, download=False)\n",
    "        else:\n",
    "            train_ds = dl_obj(datadir,clientDataTrain[clientNum],clientLabelTrain[clientNum], \n",
    "                          train=True,download=False)\n",
    "            test_ds = dl_obj(datadir,clientDataTest[clientNum],clientLabelTest[clientNum],train=False, download=False)\n",
    "        train_dl = data.DataLoader(\n",
    "            dataset=train_ds, batch_size=train_bs, shuffle=True)\n",
    "        test_dl = data.DataLoader(\n",
    "            dataset=test_ds, batch_size=test_bs, shuffle=False)\n",
    "    return train_dl, test_dl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_net(net_id, net, train_dataloader, test_dataloader, epochs, lr, args, device=\"cpu\"):\n",
    "    logger.info('Training network %s' % str(net_id))\n",
    "    logger.info('n_training: %d' % len(train_dataloader))\n",
    "    logger.info('n_test: %d' % len(test_dataloader))\n",
    "\n",
    "    train_acc = compute_accuracy(net, train_dataloader, device=device)\n",
    "    test_acc, conf_matrix = compute_accuracy(\n",
    "        net, test_dataloader, get_confusion_matrix=True, device=device)\n",
    "\n",
    "    logger.info('>> Pre-Training Training accuracy: {}'.format(train_acc))\n",
    "    logger.info('>> Pre-Training Test accuracy: {}'.format(test_acc))\n",
    "\n",
    "    if args.dataset == \"cinic10\":\n",
    "        optimizer = optim.SGD(net.parameters(), lr=lr,\n",
    "                              momentum=0.9, weight_decay=0.0001)\n",
    "        scheduler = torch.optim.lr_scheduler.StepLR(\n",
    "            optimizer, step_size=1, gamma=0.95)\n",
    "    else:\n",
    "        optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9)\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss().to(device)\n",
    "\n",
    "    cnt = 0\n",
    "    losses, running_losses = [], []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        epoch_loss_collector = []\n",
    "        for batch_idx, (x, target) in enumerate(train_dataloader):\n",
    "            x, target = x.to(device), target.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            x.requires_grad = True\n",
    "            target.requires_grad = False\n",
    "            target = target.long()\n",
    "\n",
    "            out = net(x)\n",
    "            loss = criterion(out, target)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            cnt += 1\n",
    "            epoch_loss_collector.append(loss.item())\n",
    "\n",
    "        #logging.debug('Epoch: %d Loss: %f L2 loss: %f' % (epoch, loss.item(), reg*l2_reg))\n",
    "        epoch_loss = sum(epoch_loss_collector) / len(epoch_loss_collector)\n",
    "        logger.info('Epoch: %d Loss: %f' % (epoch, epoch_loss))\n",
    "\n",
    "        if args.dataset == \"cinic10\":\n",
    "            scheduler.step()\n",
    "\n",
    "    train_acc = compute_accuracy(net, train_dataloader, device=device)\n",
    "    test_acc, conf_matrix = compute_accuracy(\n",
    "        net, test_dataloader, get_confusion_matrix=True, device=device)\n",
    "\n",
    "    logger.info('>> Training accuracy: %f' % train_acc)\n",
    "    logger.info('>> Test accuracy: %f' % test_acc)\n",
    "    logger.info(' ** Training complete **')\n",
    "    return train_acc, test_acc\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_train(nets, args, net_dataidx_map, device=\"cpu\"):\n",
    "    # save local dataset\n",
    "    local_datasets = []\n",
    "    for net_id, net in nets.items():\n",
    "        if args.retrain:\n",
    "            dataidxs = net_dataidx_map[net_id]\n",
    "            # move the model to cuda device:\n",
    "            net.to(device)\n",
    "\n",
    "            train_dl_local, test_dl_local = get_dataloader(\n",
    "                args.dataset, args_datadir, args.batch_size, 32, dataidxs,net_id)\n",
    "\n",
    "            local_datasets.append((train_dl_local, test_dl_local))\n",
    "\n",
    "            # switch to global test set here\n",
    "            trainacc, testacc = train_net(\n",
    "                net_id, net, train_dl_local, test_dl_local, args.epochs, args.lr, args, device=device)\n",
    "            # saving the trained models here\n",
    "    nets_list = list(nets.values())\n",
    "    return nets_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "# Assuming that we are on a CUDA machine, this should print a CUDA device:\n",
    "logger.info(\"torch.cuda.is_available: {} device: {}\".format(torch.cuda.is_available(),device))\n",
    "\n",
    "seed = 0\n",
    "\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_models( n_nets, args):\n",
    "    '''\n",
    "    Initialize the local LeNets\n",
    "    Please note that this part is hard coded right now\n",
    "    '''\n",
    "\n",
    "    cnns = {net_i: None for net_i in range(n_nets)}\n",
    "\n",
    "    # we add this book keeping to store meta data of model weights\n",
    "    model_meta_data = []\n",
    "    layer_type = []\n",
    "\n",
    "    for cnn_i in range(n_nets):\n",
    "        if args.model == \"lenet\":\n",
    "            cnn = LeNet()\n",
    "        elif args.model == \"vgg\":\n",
    "            cnn = vgg11()\n",
    "        elif args.model == \"simple-cnn\":\n",
    "            if args.dataset in (\"cifar10\", \"cinic10\"):\n",
    "                cnn = SimpleCNN(input_dim=(92 * kernelSize),\n",
    "                                hidden_dims=[120, 84], output_dim=activityCount)\n",
    "        cnns[cnn_i] = cnn\n",
    "\n",
    "    for (k, v) in cnns[0].state_dict().items():\n",
    "        model_meta_data.append(v.shape)\n",
    "        layer_type.append(k)\n",
    "    return cnns, model_meta_data, layer_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dicts in traindata_cls_counts: \n",
    "    for keys in traindata_cls_counts[dicts]: \n",
    "        traindata_cls_counts[dicts][keys] = np.int64(traindata_cls_counts[dicts][keys])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_classes = len(np.unique(y_train))\n",
    "averaging_weights = np.zeros((args.n_nets, n_classes), dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coefficient per class\n",
    "for i in range(n_classes):\n",
    "    total_num_counts = 0\n",
    "    worker_class_counts = [0] * args.n_nets\n",
    "    for j in range(args.n_nets):\n",
    "        if i in traindata_cls_counts[j].keys():\n",
    "            total_num_counts += traindata_cls_counts[j][i]\n",
    "            worker_class_counts[j] = traindata_cls_counts[j][i]\n",
    "        else:\n",
    "            total_num_counts += 0\n",
    "            worker_class_counts[j] = 0\n",
    "    averaging_weights[:, i] = worker_class_counts / total_num_counts\n",
    "logger.info(\"averaging_weights: {}\".format(averaging_weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def roundNumber(toRoundNb):\n",
    "    return round(np.mean(toRoundNb), 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CIFAR10_truncated(data.Dataset):\n",
    "\n",
    "    def __init__(self, root,trainData,trainLabel, dataidxs=None, train=True, transform=None, target_transform=None, download=False):\n",
    "\n",
    "        self.root = root\n",
    "        self.dataidxs = dataidxs\n",
    "        self.train = train\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "        self.download = download\n",
    "\n",
    "        if self.dataidxs is not None:\n",
    "            self.data = trainData[self.dataidxs]\n",
    "            self.target = trainLabel[self.dataidxs]\n",
    "        else:\n",
    "            self.data = trainData\n",
    "            self.target = trainLabel\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img, target = self.data[index], self.target[index]\n",
    "\n",
    "\n",
    "        return img, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_accuracy(model, dataloader, get_confusion_matrix=False, device=\"cpu\"):\n",
    "\n",
    "    was_training = False\n",
    "    if model.training:\n",
    "        model.eval()\n",
    "        was_training = True\n",
    "\n",
    "    true_labels_list, pred_labels_list = np.array([]), np.array([])\n",
    "    correct, total = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (x, target) in enumerate(dataloader):\n",
    "#             print(target.shape)\n",
    "            x, target = x.to(device), target.to(device)\n",
    "            out = model(x)\n",
    "            _, pred_label = torch.max(out.data, 1)\n",
    "            \n",
    "            total += x.data.size()[0]\n",
    "\n",
    "            correct += (pred_label == target.data).sum().item()\n",
    "            if device == \"cpu\":\n",
    "                pred_labels_list = np.append(\n",
    "                    pred_labels_list, pred_label.numpy())\n",
    "                true_labels_list = np.append(\n",
    "                    true_labels_list, target.data.numpy())\n",
    "            else:\n",
    "                pred_labels_list = np.append(\n",
    "                    pred_labels_list, pred_label.cpu().numpy())\n",
    "                true_labels_list = np.append(\n",
    "                    true_labels_list, target.data.cpu().numpy())\n",
    "\n",
    "    if get_confusion_matrix:\n",
    "        conf_matrix = confusion_matrix(true_labels_list, pred_labels_list)\n",
    "\n",
    "    if was_training:\n",
    "        model.train()\n",
    "\n",
    "    if get_confusion_matrix:\n",
    "        return correct/float(total), conf_matrix\n",
    "\n",
    "    return correct/float(total)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_weights(weights):\n",
    "    Z = np.array([])\n",
    "    eps = 1e-6\n",
    "    weights_norm = {}\n",
    "    for _, weight in weights.items():\n",
    "        if len(Z) == 0:\n",
    "            Z = weight.data.numpy()\n",
    "        else:\n",
    "            Z = Z + weight.data.numpy()\n",
    "    for mi, weight in weights.items():\n",
    "        weights_norm[mi] = weight / torch.from_numpy(Z + eps)\n",
    "    return weights_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_weighted_average_pred(models: list, weights: dict, x, device=\"cpu\"):\n",
    "    out_weighted = None\n",
    "    # Compute the predictions\n",
    "    for model_i, model in enumerate(models):\n",
    "        out = F.softmax(model(x), dim=-1)  # (N, C)\n",
    "\n",
    "        weight = weights[model_i].to(device)\n",
    "        if out_weighted is None:\n",
    "            weight = weight.to(device)\n",
    "            out_weighted = (out * weight)\n",
    "        else:\n",
    "            out_weighted += (out * weight)\n",
    "\n",
    "    return out_weighted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_ensemble_accuracy(models: list, dataloader, n_classes, train_cls_counts=None, uniform_weights=False, sanity_weights=False, device=\"cpu\"):\n",
    "\n",
    "    correct, total = 0, 0\n",
    "    true_labels_list, pred_labels_list = np.array([]), np.array([])\n",
    "\n",
    "    was_training = [False]*len(models)\n",
    "    for i, model in enumerate(models):\n",
    "        if model.training:\n",
    "            was_training[i] = True\n",
    "            model.eval()\n",
    "\n",
    "    if uniform_weights is True:\n",
    "        weights_list = prepare_uniform_weights(n_classes, len(models))\n",
    "    elif sanity_weights is True:\n",
    "        weights_list = prepare_sanity_weights(n_classes, len(models))\n",
    "    else:\n",
    "        weights_list = prepare_weight_matrix(n_classes, train_cls_counts)\n",
    "    weights_norm = normalize_weights(weights_list)\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (x, target) in enumerate(dataloader):\n",
    "            x, target = x.to(device), target.to(device)\n",
    "            target = target.long()\n",
    "            out = get_weighted_average_pred(\n",
    "                models, weights_norm, x, device=device)\n",
    "\n",
    "            _, pred_label = torch.max(out, 1)\n",
    "\n",
    "            total += x.data.size()[0]\n",
    "            correct += (pred_label == target.data).sum().item()\n",
    "\n",
    "            if device == \"cpu\":\n",
    "                pred_labels_list = np.append(\n",
    "                    pred_labels_list, pred_label.numpy())\n",
    "                true_labels_list = np.append(\n",
    "                    true_labels_list, target.data.numpy())\n",
    "            else:\n",
    "                pred_labels_list = np.append(\n",
    "                    pred_labels_list, pred_label.cpu().numpy())\n",
    "                true_labels_list = np.append(\n",
    "                    true_labels_list, target.data.cpu().numpy())\n",
    "\n",
    "    conf_matrix = confusion_matrix(true_labels_list, pred_labels_list)\n",
    "\n",
    "    for i, model in enumerate(models):\n",
    "        if was_training[i]:\n",
    "            model.train()\n",
    "\n",
    "    return correct / float(total), conf_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def row_param_cost(global_weights, weights_j_l, global_sigmas, sigma_inv_j):\n",
    "\n",
    "    match_norms = ((weights_j_l + global_weights) ** 2 / (sigma_inv_j + global_sigmas)).sum(axis=1) - (\n",
    "        global_weights ** 2 / global_sigmas).sum(axis=1)\n",
    "\n",
    "    return match_norms\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_softmax_bias(batch_weights, last_layer_const, sigma, sigma0):\n",
    "    J = len(batch_weights)\n",
    "    sigma_bias = sigma\n",
    "    sigma0_bias = sigma0\n",
    "    mu0_bias = 0.1\n",
    "    softmax_bias = [batch_weights[j][-1] for j in range(J)]\n",
    "    softmax_inv_sigma = [s / sigma_bias for s in last_layer_const]\n",
    "    softmax_bias = sum([b * s for b, s in zip(softmax_bias, softmax_inv_sigma)]) + mu0_bias / sigma0_bias\n",
    "    softmax_inv_sigma = 1 / sigma0_bias + sum(softmax_inv_sigma)\n",
    "    return softmax_bias, softmax_inv_sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def row_param_cost_simplified(global_weights, weights_j_l, sij_p_gs, red_term):\n",
    "    match_norms = ((weights_j_l + global_weights) ** 2 / sij_p_gs).sum(axis=1) - red_term\n",
    "    return match_norms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_cost(global_weights, weights_j, global_sigmas, sigma_inv_j, prior_mean_norm, prior_inv_sigma,\n",
    "                 popularity_counts, gamma, J):\n",
    "\n",
    "    param_cost_start = time.time()\n",
    "    Lj = weights_j.shape[0]\n",
    "    counts = np.minimum(np.array(popularity_counts, dtype=np.float32), 10)\n",
    "\n",
    "    sij_p_gs = sigma_inv_j + global_sigmas\n",
    "    red_term = (global_weights ** 2 / global_sigmas).sum(axis=1)\n",
    "    stupid_line_start = time.time()\n",
    "\n",
    "    param_cost = np.array([row_param_cost_simplified(global_weights, weights_j[l], sij_p_gs, red_term) for l in range(Lj)], dtype=np.float32)\n",
    "    stupid_line_dur = time.time() - stupid_line_start\n",
    "\n",
    "    param_cost += np.log(counts / (J - counts))\n",
    "    param_cost_dur = time.time() - param_cost_start\n",
    "\n",
    "    ## Nonparametric cost\n",
    "    nonparam_start = time.time()\n",
    "    L = global_weights.shape[0]\n",
    "    max_added = min(Lj, max(700 - L, 1))\n",
    "    nonparam_cost = np.outer((((weights_j + prior_mean_norm) ** 2 / (prior_inv_sigma + sigma_inv_j)).sum(axis=1) - (\n",
    "                prior_mean_norm ** 2 / prior_inv_sigma).sum()), np.ones(max_added, dtype=np.float32))\n",
    "    cost_pois = 2 * np.log(np.arange(1, max_added + 1))\n",
    "    nonparam_cost -= cost_pois\n",
    "    nonparam_cost += 2 * np.log(gamma / J)\n",
    "\n",
    "    nonparam_dur = time.time() - nonparam_start\n",
    "\n",
    "    full_cost = np.hstack((param_cost, nonparam_cost)).astype(np.float32)\n",
    "    return full_cost\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reconstruct_local_net(weights, args, ori_assignments=None, worker_index=0):\n",
    "    if args.model == \"simple-cnn\":\n",
    "        if args.dataset in (\"cifar10\", \"cinic10\"):\n",
    "            input_channel = 6\n",
    "\n",
    "        num_filters = [weights[0].shape[0]]\n",
    "        input_dim = weights[2].shape[0]\n",
    "        hidden_dims = [weights[2].shape[1]]\n",
    "        matched_cnn = SimpleCNNContainer(input_channel=input_channel,\n",
    "                                         num_filters=num_filters,\n",
    "                                         kernel_size=kernelSize,\n",
    "                                         input_dim=input_dim,\n",
    "                                         hidden_dims=hidden_dims,\n",
    "                                         output_dim=activityCount)\n",
    "        shape_estimator = SimpleCNNContainerConvBlocks(input_channel=6, num_filters=num_filters, kernel_size=kernelSize,output_dim=activityCount)\n",
    "        dummy_input = torch.rand(1,6,128)\n",
    "        estimated_output = shape_estimator(dummy_input)\n",
    "        input_dim = estimated_output.view(-1).size()[0]\n",
    "\n",
    "    def __reconstruct_weights(weight, assignment, layer_ori_shape, matched_num_filters=None, weight_type=\"conv_weight\", slice_dim=\"filter\"):\n",
    "        if weight_type == \"conv_weight\":\n",
    "            if slice_dim == \"filter\":\n",
    "                res_weight = weight[assignment, :]\n",
    "            elif slice_dim == \"channel\":\n",
    "                _ori_matched_shape = list(copy.deepcopy(layer_ori_shape))\n",
    "                _ori_matched_shape[1] = matched_num_filters\n",
    "                trans_weight = trans_next_conv_layer_forward(\n",
    "                    weight, _ori_matched_shape)\n",
    "                logger.info(\"trans_weight{} assignment {} layer_ori_shape {}\".format(\n",
    "                    np.asarray(trans_weight).shape, np.asarray(assignment).shape, layer_ori_shape))\n",
    "                sliced_weight = trans_weight[assignment, :]\n",
    "                res_weight = trans_next_conv_layer_backward(\n",
    "                    sliced_weight, layer_ori_shape)\n",
    "        elif weight_type == \"bias\":\n",
    "            res_weight = weight[assignment]\n",
    "        elif weight_type == \"first_fc_weight\":\n",
    "            # NOTE: please note that in this case, we pass the `estimated_shape` to `layer_ori_shape`:\n",
    "            __ori_shape = weight.shape\n",
    "            res_weight = weight.reshape(\n",
    "                matched_num_filters, layer_ori_shape[2]*__ori_shape[1])[assignment, :]\n",
    "            res_weight = res_weight.reshape(\n",
    "                (len(assignment)*layer_ori_shape[2], __ori_shape[1]))\n",
    "        elif weight_type == \"fc_weight\":\n",
    "            if slice_dim == \"filter\":\n",
    "                res_weight = weight.T[assignment, :]\n",
    "                #res_weight = res_weight.T\n",
    "            elif slice_dim == \"channel\":\n",
    "                res_weight = weight[assignment, :]\n",
    "        return res_weight\n",
    "\n",
    "    reconstructed_weights = []\n",
    "    # handle the conv layers part which is not changing\n",
    "    for param_idx, (key_name, param) in enumerate(matched_cnn.state_dict().items()):\n",
    "        _matched_weight = weights[param_idx]\n",
    "        if param_idx < 1:  # we need to handle the 1st conv layer specificly since the color channels are aligned\n",
    "            _assignment = ori_assignments[int(param_idx / 2)][worker_index]\n",
    "            _res_weight = __reconstruct_weights(weight=_matched_weight, assignment=_assignment,\n",
    "                                                layer_ori_shape=param.size(), matched_num_filters=None,\n",
    "                                                weight_type=\"conv_weight\", slice_dim=\"filter\")\n",
    "            reconstructed_weights.append(_res_weight)\n",
    "        elif (param_idx >= 1) and (param_idx < len(weights) - 2):\n",
    "            if \"bias\" in key_name:  # the last bias layer is already aligned so we won't need to process it\n",
    "                _assignment = ori_assignments[int(param_idx / 2)][worker_index]\n",
    "                _res_bias = __reconstruct_weights(weight=_matched_weight, assignment=_assignment,\n",
    "                                                  layer_ori_shape=param.size(), matched_num_filters=None,\n",
    "                                                  weight_type=\"bias\", slice_dim=None)\n",
    "                reconstructed_weights.append(_res_bias)\n",
    "\n",
    "            elif \"conv\" in key_name or \"features\" in key_name:\n",
    "                # we make a note here that for these weights, we will need to slice in both `filter` and `channel` dimensions\n",
    "                cur_assignment = ori_assignments[int(\n",
    "                    param_idx / 2)][worker_index]\n",
    "                prev_assignment = ori_assignments[int(\n",
    "                    param_idx / 2)-1][worker_index]\n",
    "                _matched_num_filters = weights[param_idx - 2].shape[0]\n",
    "                _layer_ori_shape = list(param.size())\n",
    "                _layer_ori_shape[0] = _matched_weight.shape[0]\n",
    "\n",
    "                _temp_res_weight = __reconstruct_weights(weight=_matched_weight, assignment=prev_assignment,\n",
    "                                                         layer_ori_shape=_layer_ori_shape, matched_num_filters=_matched_num_filters,\n",
    "                                                         weight_type=\"conv_weight\", slice_dim=\"channel\")\n",
    "\n",
    "                _res_weight = __reconstruct_weights(weight=_temp_res_weight, assignment=cur_assignment,\n",
    "                                                    layer_ori_shape=param.size(), matched_num_filters=None,\n",
    "                                                    weight_type=\"conv_weight\", slice_dim=\"filter\")\n",
    "                reconstructed_weights.append(_res_weight)\n",
    "\n",
    "            elif \"fc\" in key_name or \"classifier\" in key_name:\n",
    "                # we make a note here that for these weights, we will need to slice in both `filter` and `channel` dimensions\n",
    "                cur_assignment = ori_assignments[int(\n",
    "                    param_idx / 2)][worker_index]\n",
    "                prev_assignment = ori_assignments[int(\n",
    "                    param_idx / 2)-1][worker_index]\n",
    "                _matched_num_filters = weights[param_idx - 2].shape[0]\n",
    "\n",
    "                if param_idx != 2:  # this is the index of the first fc layer\n",
    "                    #logger.info(\"%%%%%%%%%%%%%%% prev assignment length: {}, cur assignmnet length: {}\".format(len(prev_assignment), len(cur_assignment)))\n",
    "                    temp_res_weight = __reconstruct_weights(weight=_matched_weight, assignment=prev_assignment,\n",
    "                                                            layer_ori_shape=param.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                            weight_type=\"fc_weight\", slice_dim=\"channel\")\n",
    "\n",
    "                    _res_weight = __reconstruct_weights(weight=temp_res_weight, assignment=cur_assignment,\n",
    "                                                        layer_ori_shape=param.size(), matched_num_filters=None,\n",
    "                                                        weight_type=\"fc_weight\", slice_dim=\"filter\")\n",
    "\n",
    "                    reconstructed_weights.append(_res_weight.T)\n",
    "                else:\n",
    "                    # that's for handling the first fc layer that is connected to the conv blocks\n",
    "                    temp_res_weight = __reconstruct_weights(weight=_matched_weight, assignment=prev_assignment,\n",
    "                                                            layer_ori_shape=estimated_output.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                            weight_type=\"first_fc_weight\", slice_dim=None)\n",
    "\n",
    "                    _res_weight = __reconstruct_weights(weight=temp_res_weight, assignment=cur_assignment,\n",
    "                                                        layer_ori_shape=param.size(), matched_num_filters=None,\n",
    "                                                        weight_type=\"fc_weight\", slice_dim=\"filter\")\n",
    "                    reconstructed_weights.append(_res_weight.T)\n",
    "        elif param_idx == len(weights) - 2:\n",
    "            # this is to handle the weight of the last layer\n",
    "            prev_assignment = ori_assignments[int(\n",
    "                param_idx / 2)-1][worker_index]\n",
    "            _res_weight = _matched_weight[prev_assignment, :]\n",
    "            reconstructed_weights.append(_res_weight)\n",
    "        elif param_idx == len(weights) - 1:\n",
    "            reconstructed_weights.append(_matched_weight)\n",
    "\n",
    "    return reconstructed_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def matching_upd_j(weights_j, global_weights, sigma_inv_j, global_sigmas, prior_mean_norm, prior_inv_sigma,\n",
    "                   popularity_counts, gamma, J):\n",
    "\n",
    "    L = global_weights.shape[0]\n",
    "\n",
    "    compute_cost_start = time.time()\n",
    "    full_cost = compute_cost(global_weights.astype(np.float32), weights_j.astype(np.float32), global_sigmas.astype(np.float32), sigma_inv_j.astype(np.float32), prior_mean_norm.astype(np.float32), prior_inv_sigma.astype(np.float32),\n",
    "                             popularity_counts, gamma, J)\n",
    "    compute_cost_dur = time.time() - compute_cost_start\n",
    "    #logger.info(\"###### Compute cost dur: {}\".format(compute_cost_dur))\n",
    "\n",
    "    #row_ind, col_ind = linear_sum_assignment(-full_cost)\n",
    "    # please note that this can not run on non-Linux systems\n",
    "    start_time = time.time()\n",
    "    row_ind, col_ind = solve_dense(-full_cost)\n",
    "    solve_dur = time.time() - start_time\n",
    "\n",
    "\n",
    "\n",
    "    assignment_j = []\n",
    "\n",
    "    new_L = L\n",
    "\n",
    "    for l, i in zip(row_ind, col_ind):\n",
    "        if i < L:\n",
    "            popularity_counts[i] += 1\n",
    "            assignment_j.append(i)\n",
    "            global_weights[i] += weights_j[l]\n",
    "            global_sigmas[i] += sigma_inv_j\n",
    "        else:  # new neuron\n",
    "            popularity_counts += [1]\n",
    "            assignment_j.append(new_L)\n",
    "            new_L += 1\n",
    "            global_weights = np.vstack((global_weights, prior_mean_norm + weights_j[l]))\n",
    "            global_sigmas = np.vstack((global_sigmas, prior_inv_sigma + sigma_inv_j))\n",
    "\n",
    "    return global_weights, global_sigmas, popularity_counts, assignment_j\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdm_prepare_full_weights_cnn(nets, device=\"cpu\"):\n",
    "    \"\"\"\n",
    "    we extract all weights of the conv nets out here:\n",
    "    \"\"\"\n",
    "    weights = []\n",
    "    for net_i, net in enumerate(nets):\n",
    "        net_weights = []\n",
    "        statedict = net.state_dict()\n",
    "\n",
    "        for param_id, (k, v) in enumerate(statedict.items()):\n",
    "            if device == \"cpu\":\n",
    "                if 'fc' in k or 'classifier' in k:\n",
    "                    if 'weight' in k:\n",
    "                        net_weights.append(v.numpy().T)\n",
    "                    else:\n",
    "                        net_weights.append(v.numpy())\n",
    "                elif 'conv' in k or 'features' in k:\n",
    "                    if 'weight' in k:\n",
    "                        _weight_shape = v.size()\n",
    "                        if len(_weight_shape) == 3:\n",
    "                            net_weights.append(v.numpy().reshape(_weight_shape[0], _weight_shape[1]*_weight_shape[2]))\n",
    "                        else:\n",
    "                            pass\n",
    "                    else:\n",
    "                        net_weights.append(v.numpy())\n",
    "            else:\n",
    "                if 'fc' in k or 'classifier' in k:\n",
    "                    if 'weight' in k:\n",
    "                        net_weights.append(v.cpu().numpy().T)\n",
    "                    else:\n",
    "                        net_weights.append(v.cpu().numpy())\n",
    "                elif 'conv' in k or 'features' in k:\n",
    "                    if 'weight' in k:\n",
    "                        _weight_shape = v.size()\n",
    "                        if len(_weight_shape) == 3:\n",
    "                            net_weights.append(v.cpu().numpy().reshape(_weight_shape[0], _weight_shape[1]*_weight_shape[2]))\n",
    "                        else:\n",
    "                            pass\n",
    "                    else:\n",
    "                        net_weights.append(v.cpu().numpy())\n",
    "        weights.append(net_weights)\n",
    "    return weights\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def block_patching(w_j, L_next, assignment_j_c, layer_index, model_meta_data, \n",
    "                                matching_shapes=None, \n",
    "                                layer_type=\"fc\", \n",
    "                                dataset=\"cifar10\",\n",
    "                                network_name=\"lenet\"):\n",
    "    \"\"\"\n",
    "    In CNN, weights patching needs to be handled block-wisely\n",
    "    We handle all conv layers and the first fc layer connected with the output of conv layers here\n",
    "    \"\"\"\n",
    "\n",
    "    if assignment_j_c is None:\n",
    "        return w_j\n",
    "\n",
    "    layer_meta_data = model_meta_data[2 * layer_index - 2]\n",
    "    prev_layer_meta_data = model_meta_data[2 * layer_index - 2 - 2]\n",
    "\n",
    "    if layer_type == \"conv\":    \n",
    "        new_w_j = np.zeros((w_j.shape[0], L_next*(layer_meta_data[-1])))\n",
    "\n",
    "        # we generate a sequence of block indices\n",
    "        block_indices = [np.arange(i*layer_meta_data[-1], (i+1)*layer_meta_data[-1]) for i in range(L_next)]\n",
    "        ori_block_indices = [np.arange(i*layer_meta_data[-1], (i+1)*layer_meta_data[-1]) for i in range(layer_meta_data[1])]\n",
    "        for ori_id in range(layer_meta_data[1]):\n",
    "            new_w_j[:, block_indices[assignment_j_c[ori_id]]] = w_j[:, ori_block_indices[ori_id]]\n",
    "\n",
    "    elif layer_type == \"fc\":\n",
    "        # we need to estimate the output shape here:\n",
    "        if network_name == \"simple-cnn\":\n",
    "            if dataset in (\"cifar10\", \"cinic10\"):\n",
    "                shape_estimator = SimpleCNNContainerConvBlocks(input_channel=6, num_filters=matching_shapes, kernel_size=kernelSize,output_dim=activityCount)\n",
    "        if dataset in (\"cifar10\", \"cinic10\"):\n",
    "            dummy_input = torch.rand(1, 6, 128)\n",
    "        estimated_output = shape_estimator(dummy_input)\n",
    "        new_w_j = np.zeros((w_j.shape[0], estimated_output.view(-1).size()[0]))\n",
    "#         logger.info(\"estimated_output shape : {}\".format(estimated_output.size()))\n",
    "#         logger.info(\"meta data of previous layer: {}\".format(prev_layer_meta_data))\n",
    "        \n",
    "        block_indices = [np.arange(i*estimated_output.size()[-1], (i+1)*estimated_output.size()[-1]) for i in range(L_next)]\n",
    "\n",
    "        ori_block_indices = [np.arange(i*estimated_output.size()[-1], (i+1)*estimated_output.size()[-1]) for i in range(prev_layer_meta_data[0])]\n",
    "\n",
    "        for ori_id in range(prev_layer_meta_data[0]):\n",
    "            #logger.info(\"{} ------------ to ------------ {}\".format(block_indices[assignment_j_c[ori_id]], ori_block_indices[ori_id]))\n",
    "            new_w_j[:, block_indices[assignment_j_c[ori_id]]] = w_j[:, ori_block_indices[ori_id]]\n",
    "    return new_w_j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def match_layer(weights_bias, sigma_inv_layer, mean_prior, sigma_inv_prior, gamma, it):\n",
    "    J = len(weights_bias)\n",
    "\n",
    "    group_order = sorted(range(J), key=lambda x: -weights_bias[x].shape[0])\n",
    "\n",
    "    batch_weights_norm = [w * s for w, s in zip(weights_bias, sigma_inv_layer)]\n",
    "    prior_mean_norm = mean_prior * sigma_inv_prior\n",
    "\n",
    "    global_weights = prior_mean_norm + batch_weights_norm[group_order[0]]\n",
    "    global_sigmas = np.outer(np.ones(global_weights.shape[0]), sigma_inv_prior + sigma_inv_layer[group_order[0]])\n",
    "\n",
    "    popularity_counts = [1] * global_weights.shape[0]\n",
    "\n",
    "    assignment = [[] for _ in range(J)]\n",
    "\n",
    "    assignment[group_order[0]] = list(range(global_weights.shape[0]))\n",
    "\n",
    "    ## Initialize\n",
    "    for j in group_order[1:]:\n",
    "        global_weights, global_sigmas, popularity_counts, assignment_j = matching_upd_j(batch_weights_norm[j],\n",
    "                                                                                        global_weights,\n",
    "                                                                                        sigma_inv_layer[j],\n",
    "                                                                                        global_sigmas, prior_mean_norm,\n",
    "                                                                                        sigma_inv_prior,\n",
    "                                                                                        popularity_counts, gamma, J)\n",
    "        assignment[j] = assignment_j\n",
    "\n",
    "    ## Iterate over groups\n",
    "    for iteration in range(it):\n",
    "        random_order = np.random.permutation(J)\n",
    "        for j in random_order:  # random_order:\n",
    "            to_delete = []\n",
    "            ## Remove j\n",
    "            Lj = len(assignment[j])\n",
    "            for l, i in sorted(zip(range(Lj), assignment[j]), key=lambda x: -x[1]):\n",
    "                popularity_counts[i] -= 1\n",
    "                if popularity_counts[i] == 0:\n",
    "                    del popularity_counts[i]\n",
    "                    to_delete.append(i)\n",
    "                    for j_clean in range(J):\n",
    "                        for idx, l_ind in enumerate(assignment[j_clean]):\n",
    "                            if i < l_ind and j_clean != j:\n",
    "                                assignment[j_clean][idx] -= 1\n",
    "                            elif i == l_ind and j_clean != j:\n",
    "                                logger.info('Warning - weird unmatching')\n",
    "                else:\n",
    "                    global_weights[i] = global_weights[i] - batch_weights_norm[j][l]\n",
    "                    global_sigmas[i] -= sigma_inv_layer[j]\n",
    "\n",
    "            global_weights = np.delete(global_weights, to_delete, axis=0)\n",
    "            global_sigmas = np.delete(global_sigmas, to_delete, axis=0)\n",
    "\n",
    "            ## Match j\n",
    "            global_weights, global_sigmas, popularity_counts, assignment_j = matching_upd_j(batch_weights_norm[j],\n",
    "                                                                                            global_weights,\n",
    "                                                                                            sigma_inv_layer[j],\n",
    "                                                                                            global_sigmas,\n",
    "                                                                                            prior_mean_norm,\n",
    "                                                                                            sigma_inv_prior,\n",
    "                                                                                            popularity_counts, gamma, J)\n",
    "            assignment[j] = assignment_j\n",
    "\n",
    "    logger.info('Number of global neurons is %d, gamma %f' % (global_weights.shape[0], gamma))\n",
    "    logger.info(\"***************Shape of global weights after match: {} ******************\".format(global_weights.shape))\n",
    "    return assignment, global_weights, global_sigmas\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_wise_group_descent(batch_weights, layer_index, batch_frequencies, sigma_layers, \n",
    "                                sigma0_layers, gamma_layers, it, \n",
    "                                model_meta_data, \n",
    "                                model_layer_type,\n",
    "                                n_layers,\n",
    "                                matching_shapes,\n",
    "                                args):\n",
    "    \"\"\"\n",
    "    We implement a layer-wise matching here:\n",
    "    \"\"\"\n",
    "    if type(sigma_layers) is not list:\n",
    "        sigma_layers = (n_layers - 1) * [sigma_layers]\n",
    "    if type(sigma0_layers) is not list:\n",
    "        sigma0_layers = (n_layers - 1) * [sigma0_layers]\n",
    "    if type(gamma_layers) is not list:\n",
    "        gamma_layers = (n_layers - 1) * [gamma_layers]\n",
    "\n",
    "    last_layer_const = []\n",
    "    total_freq = sum(batch_frequencies)\n",
    "    for f in batch_frequencies:\n",
    "        last_layer_const.append(f / total_freq)\n",
    "\n",
    "    # J: number of workers\n",
    "    J = len(batch_weights)\n",
    "    init_num_kernel = batch_weights[0][0].shape[0]\n",
    "    init_channel_kernel_dims = []\n",
    "    for bw in batch_weights[0]:\n",
    "        if len(bw.shape) > 1:\n",
    "            init_channel_kernel_dims.append(bw.shape[1])\n",
    "\n",
    "    \n",
    "    sigma_bias_layers = sigma_layers\n",
    "    sigma0_bias_layers = sigma0_layers\n",
    "    mu0 = 0.\n",
    "    mu0_bias = 0.1\n",
    "    assignment_c = [None for j in range(J)]\n",
    "    L_next = None\n",
    "\n",
    "    sigma = sigma_layers[layer_index - 1]\n",
    "    sigma_bias = sigma_bias_layers[layer_index - 1]\n",
    "    gamma = gamma_layers[layer_index - 1]\n",
    "    sigma0 = sigma0_layers[layer_index - 1]\n",
    "    sigma0_bias = sigma0_bias_layers[layer_index - 1]\n",
    "\n",
    "    if layer_index <= 1:\n",
    "        weights_bias = [np.hstack((batch_weights[j][0], batch_weights[j][layer_index * 2 - 1].reshape(-1, 1))) for j in range(J)]\n",
    "\n",
    "        sigma_inv_prior = np.array(\n",
    "            init_channel_kernel_dims[layer_index - 1] * [1 / sigma0] + [1 / sigma0_bias])\n",
    "        mean_prior = np.array(init_channel_kernel_dims[layer_index - 1] * [mu0] + [mu0_bias])\n",
    "\n",
    "        # handling 2-layer neural network\n",
    "        if n_layers == 2:\n",
    "            sigma_inv_layer = [\n",
    "                np.array(D * [1 / sigma] + [1 / sigma_bias] + [y / sigma for y in last_layer_const[j]]) for j in range(J)]\n",
    "        else:\n",
    "            sigma_inv_layer = [np.array(init_channel_kernel_dims[layer_index - 1] * [1 / sigma] + [1 / sigma_bias]) for j in range(J)]\n",
    "\n",
    "    elif layer_index == (n_layers - 1) and n_layers > 2:\n",
    "        # our assumption is that this branch will consistently handle the last fc layers\n",
    "        layer_type = model_layer_type[2 * layer_index - 2]\n",
    "        prev_layer_type = model_layer_type[2 * layer_index - 2 - 2]\n",
    "        first_fc_identifier = (('fc' in layer_type or 'classifier' in layer_type) and ('conv' in prev_layer_type or 'features' in layer_type))\n",
    "\n",
    "\n",
    "        if first_fc_identifier:\n",
    "            weights_bias = [np.hstack((batch_weights[j][2 * layer_index - 2].T, \n",
    "                                        batch_weights[j][2 * layer_index - 1].reshape(-1, 1))) for j in range(J)]\n",
    "        else:\n",
    "            weights_bias = [np.hstack((batch_weights[j][2 * layer_index - 2].T, \n",
    "                                        batch_weights[j][2 * layer_index - 1].reshape(-1, 1))) for j in range(J)]\n",
    "\n",
    "\n",
    "        sigma_inv_prior = np.array([1 / sigma0_bias] + (weights_bias[0].shape[1] - 1) * [1 / sigma0])\n",
    "        mean_prior = np.array([mu0_bias] + (weights_bias[0].shape[1] - 1) * [mu0])\n",
    "        \n",
    "    \n",
    "        sigma_inv_layer = [np.array([1 / sigma_bias] + (weights_bias[j].shape[1] - 1) * [1 / sigma]) for j in range(J)]\n",
    "\n",
    "    elif (layer_index > 1 and layer_index < (n_layers - 1)):\n",
    "        layer_type = model_layer_type[2 * layer_index - 2]\n",
    "        prev_layer_type = model_layer_type[2 * layer_index - 2 - 2]\n",
    "\n",
    "        if 'conv' in layer_type or 'features' in layer_type:\n",
    "            weights_bias = [np.hstack((batch_weights[j][2 * layer_index - 2], batch_weights[j][2 * layer_index - 1].reshape(-1, 1))) for j in range(J)]\n",
    "\n",
    "        elif 'fc' in layer_type or 'classifier' in layer_type:\n",
    "            first_fc_identifier = (('fc' in layer_type or 'classifier' in layer_type) and ('conv' in prev_layer_type or 'features' in layer_type))\n",
    "            if first_fc_identifier:\n",
    "                weights_bias = [np.hstack((batch_weights[j][2 * layer_index - 2].T, batch_weights[j][2 * layer_index - 1].reshape(-1, 1))) for j in range(J)]\n",
    "            else:\n",
    "                weights_bias = [np.hstack((batch_weights[j][2 * layer_index - 2].T, batch_weights[j][2 * layer_index - 1].reshape(-1, 1))) for j in range(J)]          \n",
    "\n",
    "        sigma_inv_prior = np.array([1 / sigma0_bias] + (weights_bias[0].shape[1] - 1) * [1 / sigma0])\n",
    "        mean_prior = np.array([mu0_bias] + (weights_bias[0].shape[1] - 1) * [mu0])\n",
    "        sigma_inv_layer = [np.array([1 / sigma_bias] + (weights_bias[j].shape[1] - 1) * [1 / sigma]) for j in range(J)]\n",
    "\n",
    "    assignment_c, global_weights_c, global_sigmas_c = match_layer(weights_bias, sigma_inv_layer, mean_prior,\n",
    "                                                                  sigma_inv_prior, gamma, it)\n",
    "\n",
    "    L_next = global_weights_c.shape[0]\n",
    "\n",
    "    if layer_index <= 1:\n",
    "        if n_layers == 2:\n",
    "            softmax_bias, softmax_inv_sigma = process_softmax_bias(batch_weights, last_layer_const, sigma, sigma0)\n",
    "            global_weights_out = [softmax_bias]\n",
    "            global_inv_sigmas_out = [softmax_inv_sigma]\n",
    "        \n",
    "        global_weights_out = [global_weights_c[:, :init_channel_kernel_dims[int(layer_index/2)]], global_weights_c[:, init_channel_kernel_dims[int(layer_index/2)]]]\n",
    "        global_inv_sigmas_out = [global_sigmas_c[:, :init_channel_kernel_dims[int(layer_index/2)]], global_sigmas_c[:, init_channel_kernel_dims[int(layer_index/2)]]]\n",
    "\n",
    "\n",
    "\n",
    "    elif layer_index == (n_layers - 1) and n_layers > 2:\n",
    "        softmax_bias, softmax_inv_sigma = process_softmax_bias(batch_weights, last_layer_const, sigma, sigma0)\n",
    "\n",
    "        layer_type = model_layer_type[2 * layer_index - 2]\n",
    "        prev_layer_type = model_layer_type[2 * layer_index - 2 - 2]\n",
    "\n",
    "        first_fc_identifier = (('fc' in layer_type or 'classifier' in layer_type) and ('conv' in prev_layer_type or 'features' in layer_type))\n",
    "\n",
    "        # remove fitting the last layer\n",
    "        if first_fc_identifier:\n",
    "            \n",
    "            global_weights_out = [global_weights_c[:, 0:-1].T, \n",
    "                                    global_weights_c[:, -softmax_bias.shape[0]-1]]\n",
    "\n",
    "            global_inv_sigmas_out = [global_sigmas_c[:, 0:-1].T, \n",
    "                                        global_sigmas_c[:, -softmax_bias.shape[0]-1]]\n",
    "        else:\n",
    "            global_weights_out = [global_weights_c[:, 0:matching_shapes[layer_index - 1 - 1]].T, \n",
    "                                    global_weights_c[:, matching_shapes[layer_index - 1 - 1]]]\n",
    "\n",
    "            global_inv_sigmas_out = [global_sigmas_c[:, 0:matching_shapes[layer_index - 1 - 1]].T, \n",
    "                                        global_sigmas_c[:, matching_shapes[layer_index - 1 - 1]]]\n",
    "\n",
    "\n",
    "    elif (layer_index > 1 and layer_index < (n_layers - 1)):\n",
    "        layer_type = model_layer_type[2 * layer_index - 2]\n",
    "        gwc_shape = global_weights_c.shape\n",
    "\n",
    "        if \"conv\" in layer_type or 'features' in layer_type:\n",
    "            global_weights_out = [global_weights_c[:, 0:gwc_shape[1]-1], global_weights_c[:, gwc_shape[1]-1]]\n",
    "            global_inv_sigmas_out = [global_sigmas_c[:, 0:gwc_shape[1]-1], global_sigmas_c[:, gwc_shape[1]-1]]\n",
    "        elif \"fc\" in layer_type or 'classifier' in layer_type:\n",
    "            global_weights_out = [global_weights_c[:, 0:gwc_shape[1]-1].T, global_weights_c[:, gwc_shape[1]-1]]\n",
    "            global_inv_sigmas_out = [global_sigmas_c[:, 0:gwc_shape[1]-1].T, global_sigmas_c[:, gwc_shape[1]-1]]\n",
    "\n",
    "    map_out = [g_w / g_s for g_w, g_s in zip(global_weights_out, global_inv_sigmas_out)]\n",
    "    return map_out, assignment_c, L_next\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rebuild_net(weights):\n",
    "    num_filters = [weights[0].shape[0]]\n",
    "    input_dim = weights[2].shape[0]\n",
    "    hidden_dims = [weights[2].shape[1]]\n",
    "    matched_cnn = SimpleCNNContainer(6,num_filters,kernel_size=kernelSize,input_dim=input_dim,hidden_dims=hidden_dims,output_dim=activityCount)\n",
    "    return matched_cnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_model_averaging_accuracy(models, weights, train_dl, test_dl, n_classes, args):\n",
    "    \"\"\"An variant of fedaveraging\"\"\"\n",
    "    if args.model == \"lenet\":\n",
    "        avg_cnn = LeNet()\n",
    "    elif args.model == \"vgg\":\n",
    "        avg_cnn = vgg11()\n",
    "    elif args.model == \"simple-cnn\":\n",
    "        if args.dataset in (\"cifar10\", \"cinic10\"):\n",
    "            avg_cnn = SimpleCNN(input_dim=(16 * 5 * 5),\n",
    "                                hidden_dims=[120, 84], output_dim=activityCount)\n",
    "    elif args.model == \"moderate-cnn\":\n",
    "        if args.dataset in (\"cifar10\", \"cinic10\"):\n",
    "            avg_cnn = ModerateCNN()\n",
    "        elif args.dataset == \"mnist\":\n",
    "            avg_cnn = ModerateCNNMNIST()\n",
    "\n",
    "    new_state_dict = {}\n",
    "    model_counter = 0\n",
    "\n",
    "    # handle the conv layers part which is not changing\n",
    "    for param_idx, (key_name, param) in enumerate(avg_cnn.state_dict().items()):\n",
    "        if \"conv\" in key_name or \"features\" in key_name:\n",
    "            if \"weight\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(\n",
    "                    weights[param_idx].reshape(param.size()))}\n",
    "            elif \"bias\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(weights[param_idx])}\n",
    "        elif \"fc\" in key_name or \"classifier\" in key_name:\n",
    "            if \"weight\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(weights[param_idx].T)}\n",
    "            elif \"bias\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(weights[param_idx])}\n",
    "\n",
    "        new_state_dict.update(temp_dict)\n",
    "\n",
    "    avg_cnn.load_state_dict(new_state_dict)\n",
    "\n",
    "    # switch to eval mode:\n",
    "    avg_cnn.eval()\n",
    "    ##\n",
    "\n",
    "    correct, total = 0, 0\n",
    "    for batch_idx, (x, target) in enumerate(test_dl):\n",
    "        out_k = avg_cnn(x)\n",
    "        _, pred_label = torch.max(out_k, 1)\n",
    "        total += x.data.size()[0]\n",
    "        correct += (pred_label == target.data).sum().item()\n",
    "\n",
    "    logger.info(\n",
    "        \"Accuracy for Fed Averaging correct: {}, total: {}\".format(correct, total))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BBP_MAP(nets_list, model_meta_data, layer_type, net_dataidx_map,\n",
    "            averaging_weights, args,\n",
    "            device=\"cpu\"):\n",
    "    # starting the neural matching\n",
    "    models = nets_list\n",
    "    cls_freqs = traindata_cls_counts\n",
    "    n_classes = activityCount\n",
    "    it = iteration\n",
    "    sigma = args_pdm_sig\n",
    "    sigma0 = args_pdm_sig0\n",
    "    gamma = args_pdm_gamma\n",
    "    assignments_list = []\n",
    "\n",
    "    batch_weights = pdm_prepare_full_weights_cnn(models, device=device)\n",
    "    raw_batch_weights = copy.deepcopy(batch_weights)\n",
    "\n",
    "    logging.info(\"==\"*15)\n",
    "#     logging.info(\"Weights shapes: {}\".format(\n",
    "#         [bw.shape for bw in batch_weights[0]]))\n",
    "\n",
    "    batch_freqs = pdm_prepare_freq(cls_freqs, n_classes)\n",
    "    res = {}\n",
    "    best_test_acc, best_train_acc, best_weights, best_sigma, best_gamma, best_sigma0 = - \\\n",
    "        1, -1, None, -1, -1, -1\n",
    "\n",
    "    gamma = gammaValue\n",
    "    sigma = 1.0\n",
    "    sigma0 = 1.0\n",
    "\n",
    "    n_layers = int(len(batch_weights[0]) / 2)\n",
    "    num_workers = len(nets_list)\n",
    "    matching_shapes = []\n",
    "\n",
    "    first_fc_index = None\n",
    "\n",
    "    for layer_index in range(1, n_layers):\n",
    "        layer_hungarian_weights, assignment, L_next = layer_wise_group_descent(\n",
    "            batch_weights=batch_weights,\n",
    "            layer_index=layer_index,\n",
    "            sigma0_layers=sigma0,\n",
    "            sigma_layers=sigma,\n",
    "            batch_frequencies=batch_freqs,\n",
    "            it=it,\n",
    "            gamma_layers=gamma,\n",
    "            model_meta_data=model_meta_data,\n",
    "            model_layer_type=layer_type,\n",
    "            n_layers=n_layers,\n",
    "            matching_shapes=matching_shapes,\n",
    "            args=args\n",
    "        )\n",
    "        \n",
    "        assignments_list.append(assignment)\n",
    "\n",
    "        # iii) load weights to the model and train the whole thing\n",
    "        type_of_patched_layer = layer_type[2 * (layer_index + 1) - 2]\n",
    "        if 'conv' in type_of_patched_layer or 'features' in type_of_patched_layer:\n",
    "            l_type = \"conv\"\n",
    "        elif 'fc' in type_of_patched_layer or 'classifier' in type_of_patched_layer:\n",
    "            l_type = \"fc\"\n",
    "\n",
    "        type_of_this_layer = layer_type[2 * layer_index - 2]\n",
    "        type_of_prev_layer = layer_type[2 * layer_index - 2 - 2]\n",
    "        first_fc_identifier = (('fc' in type_of_this_layer or 'classifier' in type_of_this_layer) and (\n",
    "            'conv' in type_of_prev_layer or 'features' in type_of_this_layer))\n",
    "\n",
    "        if first_fc_identifier:\n",
    "            first_fc_index = layer_index\n",
    "\n",
    "        matching_shapes.append(L_next)\n",
    "        tempt_weights = [([batch_weights[w][i] for i in range(2 * layer_index - 2)] +\n",
    "                          copy.deepcopy(layer_hungarian_weights)) for w in range(num_workers)]\n",
    "\n",
    "        # i) permutate the next layer wrt matching result\n",
    "        for worker_index in range(num_workers):\n",
    "            if first_fc_index is None:\n",
    "                if l_type == \"conv\":\n",
    "                    patched_weight = block_patching(batch_weights[worker_index][2 * (layer_index + 1) - 2],\n",
    "                                                    L_next, assignment[worker_index],\n",
    "                                                    layer_index+1, model_meta_data,\n",
    "                                                    matching_shapes=matching_shapes, layer_type=l_type,\n",
    "                                                    dataset=args.dataset, network_name=args.model)\n",
    "                elif l_type == \"fc\":\n",
    "                    patched_weight = block_patching(batch_weights[worker_index][2 * (layer_index + 1) - 2].T,\n",
    "                                                    L_next, assignment[worker_index],\n",
    "                                                    layer_index+1, model_meta_data,\n",
    "                                                    matching_shapes=matching_shapes, layer_type=l_type,\n",
    "                                                    dataset=args.dataset, network_name=args.model).T\n",
    "\n",
    "            elif layer_index >= first_fc_index:\n",
    "                patched_weight = patch_weights(\n",
    "                    batch_weights[worker_index][2 * (layer_index + 1) - 2].T, L_next, assignment[worker_index]).T\n",
    "\n",
    "            tempt_weights[worker_index].append(patched_weight)\n",
    "        # ii) prepare the whole network weights\n",
    "        for worker_index in range(num_workers):\n",
    "            for lid in range(2 * (layer_index + 1) - 1, len(batch_weights[0])):\n",
    "                tempt_weights[worker_index].append(\n",
    "                    batch_weights[worker_index][lid])\n",
    "                \n",
    "        myModel = np.asarray(tempt_weights[0])\n",
    "\n",
    "                \n",
    "        for worker_index in range(num_workers):\n",
    "            dataidxs = net_dataidx_map[worker_index]\n",
    "            train_dl_local, test_dl_local = get_dataloader(args.dataset, args_datadir, args.batch_size, 32, dataidxs,worker_index)\n",
    "        retrained_nets = []\n",
    "        for worker_index in range(num_workers):\n",
    "            dataidxs = net_dataidx_map[worker_index]\n",
    "            train_dl_local, test_dl_local = get_dataloader(\n",
    "                args.dataset, args_datadir, args.batch_size, 32, dataidxs,worker_index)\n",
    "\n",
    "            retrained_cnn = local_retrain((train_dl_local, test_dl_local), tempt_weights[worker_index], args,\n",
    "                                          freezing_index=(2 * (layer_index + 1) - 2), device=device)\n",
    "            retrained_nets.append(retrained_cnn)\n",
    "        batch_weights = pdm_prepare_full_weights_cnn(\n",
    "            retrained_nets, device=device)\n",
    "\n",
    "\n",
    "    matched_weights = []\n",
    "    num_layers = len(batch_weights[0])\n",
    "\n",
    "    last_layer_weights_collector = []\n",
    "\n",
    "    for i in range(num_workers):\n",
    "        # firstly we combine last layer's weight and bias\n",
    "        bias_shape = batch_weights[i][-1].shape\n",
    "        last_layer_bias = batch_weights[i][-1].reshape((1, bias_shape[0]))\n",
    "        last_layer_weights = np.concatenate(\n",
    "            (batch_weights[i][-2], last_layer_bias), axis=0)\n",
    "\n",
    "        last_layer_weights_collector.append(last_layer_weights)\n",
    "\n",
    "    last_layer_weights_collector = np.array(last_layer_weights_collector)\n",
    "\n",
    "    avg_last_layer_weight = np.zeros(\n",
    "        last_layer_weights_collector[0].shape, dtype=np.float32)\n",
    "\n",
    "    for i in range(n_classes):\n",
    "        avg_weight_collector = np.zeros(\n",
    "            last_layer_weights_collector[0][:, 0].shape, dtype=np.float32)\n",
    "        for j in range(num_workers):\n",
    "            avg_weight_collector += averaging_weights[j][i] * \\\n",
    "                last_layer_weights_collector[j][:, i]\n",
    "        avg_last_layer_weight[:, i] = avg_weight_collector\n",
    "\n",
    "    #avg_last_layer_weight = np.mean(last_layer_weights_collector, axis=0)\n",
    "    for i in range(num_layers):\n",
    "        if i < (num_layers - 2):\n",
    "            matched_weights.append(batch_weights[0][i])\n",
    "\n",
    "    matched_weights.append(avg_last_layer_weight[0:-1, :])\n",
    "    matched_weights.append(avg_last_layer_weight[-1, :])\n",
    "    return matched_weights, assignments_list\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def computeWeights(modelWeight):\n",
    "    modelWeight = np.asarray(modelWeight)\n",
    "    modelWeightsPrep = []\n",
    "    for i in range(int(modelWeight.shape[0]/ 2)):\n",
    "        if(layerType[i] == 0 ):\n",
    "            weightReshaped = modelWeight[i*2]\n",
    "        else:\n",
    "            weightReshaped = modelWeight[i*2].T\n",
    "        biasReshaped = modelWeight[i*2+1].reshape(-1,1)\n",
    "        modelWeightsPrep.append(np.hstack((weightReshaped,biasReshaped)))\n",
    "\n",
    "    modelWeightsPrep = np.asarray(modelWeightsPrep)\n",
    "    return modelWeightsPrep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fedma_comm(batch_weights, model_meta_data, layer_type, net_dataidx_map,\n",
    "               averaging_weights, args,\n",
    "               train_dl_global,\n",
    "               test_dl_global,\n",
    "               assignments_list,\n",
    "               comm_round=2,\n",
    "               device=\"cpu\"):\n",
    "    \n",
    "    \n",
    "    timePerRound = 0\n",
    "    n_layers = int(len(batch_weights[0]) / 2)\n",
    "    num_workers = len(batch_weights)\n",
    "    matching_shapes = []\n",
    "    first_fc_index = None\n",
    "    gamma = gammaValue\n",
    "    sigma = 1.0\n",
    "    sigma0 = 1.0\n",
    "\n",
    "    cls_freqs = traindata_cls_counts\n",
    "    n_classes = activityCount\n",
    "    batch_freqs = pdm_prepare_freq(cls_freqs, n_classes)\n",
    "    it = iteration\n",
    "    serverAccuracyTest = []\n",
    "    serverLossTest = []\n",
    "    serverAccuracyTrain = []\n",
    "    serverLossTrain = []\n",
    "    \n",
    "    serverMacroVal_f1Train = []\n",
    "    serverMacroVal_f1Test = []\n",
    "    \n",
    "    clientTestSingleAccuracy = []\n",
    "    clientTestSingleAccuracyStd = []\n",
    "    clientTestSingleLoss = []\n",
    "    clientTestSingleLossStd = []    \n",
    "    \n",
    "    clientTrainSingleAccuracy = []\n",
    "    clientTrainSingleAccuracyStd = []\n",
    "    clientTrainSingleLoss = []\n",
    "    clientTrainSingleLossStd = []   \n",
    "    \n",
    "    clientTestAllAccuracy = []\n",
    "    clientTestAllAccuracyStd = []\n",
    "    clientTestAllLoss = []\n",
    "    clientTestAllLossStd = []    \n",
    "    \n",
    "    clientTrainAllAccuracy = []\n",
    "    clientTrainAllAccuracyStd = []\n",
    "    clientTrainAllLoss = []\n",
    "    clientTrainAllLossStd = []   \n",
    "    \n",
    "    cleintMacroVal_f1SingleTrain = []\n",
    "    clientMacroVal_f1SingleTest = []\n",
    "    \n",
    "    cleintMacroVal_f1AllTrain = []\n",
    "    clientMacroVal_f1AllTest = []\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    for cr in range(comm_round):\n",
    "        startTime = time.time()\n",
    "        logger.info(\"Entering communication round: {} ...\".format(cr))\n",
    "        retrained_nets = []\n",
    "        for worker_index in range(args.n_nets):\n",
    "            dataidxs = net_dataidx_map[worker_index]\n",
    "            train_dl_local, test_dl_local = get_dataloader(\n",
    "                args.dataset, args_datadir, args.batch_size, 32, dataidxs,worker_index)\n",
    "\n",
    "            recons_local_net = reconstruct_local_net(\n",
    "                batch_weights[worker_index], args, ori_assignments=assignments_list, worker_index=worker_index) \n",
    "                \n",
    "            retrained_cnn = local_retrain((train_dl_local, test_dl_local), recons_local_net, args,\n",
    "                                          mode=\"bottom-up\", freezing_index=0, ori_assignments=None, device=device)\n",
    "            retrained_nets.append(retrained_cnn)\n",
    "\n",
    "        server_weights = computeWeights(batch_weights[0])\n",
    "\n",
    "        batch_weights = pdm_prepare_full_weights_cnn(retrained_nets, device=device)\n",
    "        \n",
    "        tempSingleAccuracyTest = []\n",
    "        tempSingleLossTest = []\n",
    "        tempSingleAccuracyTrain = []\n",
    "        tempSingleLossTrain = []\n",
    "        \n",
    "        tempAllAccuracyTest = []\n",
    "        tempAllLossTest = []\n",
    "        tempAllAccuracyTrain = []\n",
    "        tempAllLossTrain = []\n",
    "        \n",
    "        tempSingleMacroVal_f1Train = []\n",
    "        tempSingleMacroVal_f1Test = []\n",
    "        \n",
    "        tempAllMacroVal_f1Train = []\n",
    "        tempAllMacroVal_f1Test = []\n",
    "        \n",
    "\n",
    "        for worker_index in range(args.n_nets): \n",
    "            dataidxs = net_dataidx_map[worker_index]\n",
    "            train_dl_local, test_dl_local = get_dataloader(args.dataset, args_datadir, args.batch_size, 32, dataidxs,worker_index)\n",
    "        \n",
    "            tempAccSingleTest, tempLsSingleTest,tempAccSingleTrain,tempLsSingleTrain,tempMacroVal_f1SingleTrain,tempMacroVal_f1SingleTest = compute_full_cnn_accuracy(models,batch_weights[worker_index],train_dl_local,test_dl_local,n_classes,args,clientServer=\"Client\")\n",
    "            \n",
    "            tempSingleAccuracyTest.append(tempAccSingleTest)\n",
    "            tempSingleLossTest.append(tempLsSingleTest)\n",
    "            \n",
    "            tempSingleAccuracyTrain.append(tempAccSingleTrain)\n",
    "            tempSingleLossTrain.append(tempLsSingleTrain)\n",
    "            \n",
    "            tempSingleMacroVal_f1Train.append(tempMacroVal_f1SingleTrain)\n",
    "            tempSingleMacroVal_f1Test.append(tempMacroVal_f1SingleTest)\n",
    "            \n",
    "            \n",
    "            if(ClientAllTest  == True):\n",
    "                \n",
    "                tempAccAllTest, tempLsAllTest,tempAccAllTrain,tempLsAllTrain, tpAllMacroVal_f1Train, tpAllMacroVal_f1Test= compute_full_cnn_accuracy(models,batch_weights[worker_index],train_dl_global,test_dl_global,n_classes,args,clientServer=\"Client\")\n",
    "                \n",
    "                tempAllAccuracyTest.append(tempAccAllTest)\n",
    "                tempAllLossTest.append(tempLsAllTest)\n",
    "\n",
    "                tempAllAccuracyTrain.append(tempAccAllTrain)\n",
    "                tempAllLossTrain.append(tempLsAllTrain)\n",
    "                \n",
    "                tempAllMacroVal_f1Train.append(tpAllMacroVal_f1Train)\n",
    "                tempAllMacroVal_f1Test.append(tpAllMacroVal_f1Test)\n",
    "                 \n",
    "\n",
    "        if(euclid):\n",
    "            meanServerClient = []\n",
    "            stdServerClient = []\n",
    "            serverShape = np.asarray(computeWeights(recons_local_net))\n",
    "            localMeanClientLayer = []\n",
    "            localStdClientLayer = []\n",
    "            for clientIndex in range(clientCount):\n",
    "                localMeanServerClient = []\n",
    "                localStdServerClient = []\n",
    "                localShape = np.asarray(computeWeights(batch_weights[clientIndex]))\n",
    "                for i in range(serverShape.shape[0]):\n",
    "                    newLayerDist = np.sqrt((serverShape[i] - localShape[i])**2)\n",
    "                    localMeanServerClient.append(np.mean(newLayerDist))\n",
    "                    localStdServerClient.append(np.std(newLayerDist))\n",
    "                localMeanClientLayer.append(localMeanServerClient)\n",
    "                localStdClientLayer.append(localStdServerClient)\n",
    "\n",
    "                meanServerClient.append(np.mean(localMeanServerClient))\n",
    "                stdServerClient.append(np.mean(localStdServerClient))\n",
    "\n",
    "\n",
    "    #         15 clients \n",
    "            meanHistoryDist.append(np.asarray(meanServerClient))\n",
    "            stdHistoryDist.append(np.asarray(stdServerClient))\n",
    "\n",
    "    #         per layer distance\n",
    "            meanRoundLayerHistory.append(np.mean(localMeanClientLayer,axis = 0))\n",
    "            stdRoundLayerHistory.append(np.mean(localStdClientLayer,axis=0))\n",
    "\n",
    "    #         all layer distance\n",
    "            meanRoundGeneralLayerHistory.append(np.mean(localMeanClientLayer))\n",
    "            stdRoundGeneralLayerHistory.append(np.mean(localStdClientLayer))\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "        clientTestSingleAccuracy.append(np.mean(tempSingleAccuracyTest))\n",
    "        clientTestSingleAccuracyStd.append(np.std(tempSingleAccuracyTest))\n",
    "        clientTestSingleLoss.append(np.mean(tempSingleLossTest))\n",
    "        clientTestSingleLossStd.append(np.mean(tempSingleLossTest))\n",
    "        \n",
    "        clientTrainSingleAccuracy.append(np.mean(tempSingleAccuracyTrain))\n",
    "        clientTrainSingleAccuracyStd.append(np.std(tempSingleAccuracyTrain))\n",
    "        clientTrainSingleLoss.append(np.mean(tempSingleLossTrain))\n",
    "        clientTrainSingleLossStd.append(np.mean(tempSingleLossTrain))\n",
    "        \n",
    "        cleintMacroVal_f1SingleTrain.append(np.mean(tempSingleMacroVal_f1Train))\n",
    "        clientMacroVal_f1SingleTest.append(np.mean(tempSingleMacroVal_f1Test))\n",
    "        \n",
    "        \n",
    "        logger.info(\"Client single train accuracy: {}\".format(np.mean(clientTrainSingleAccuracy)))\n",
    "#         logger.info(\"Client single train loss: {}\".format(np.mean(clientTrainSingleLoss)))\n",
    "        logger.info(\"Client single test accuracy: {}\".format(np.mean(clientTestSingleAccuracy)))\n",
    "#         logger.info(\"Client single test loss: {}\".format(np.mean(clientTestSingleLoss)))\n",
    "        \n",
    "        if(ClientAllTest  == True):\n",
    "            clientTestAllAccuracy.append(np.mean(tempAllAccuracyTest))\n",
    "            clientTestAllAccuracyStd.append(np.std(tempAllAccuracyTest))\n",
    "            clientTestAllLoss.append(np.mean(tempAllLossTest))\n",
    "            clientTestAllLossStd.append(np.mean(tempAllLossTest))\n",
    "\n",
    "            clientTrainAllAccuracy.append(np.mean(tempAllAccuracyTrain))\n",
    "            clientTrainAllAccuracyStd.append(np.std(tempAllAccuracyTrain))\n",
    "            clientTrainAllLoss.append(np.mean(tempAllLossTrain))\n",
    "            clientTrainAllLossStd.append(np.mean(tempAllLossTrain))\n",
    "\n",
    "            logger.info(\"Client all train accuracy: {}\".format(np.mean(clientTrainAllAccuracy)))\n",
    "#             logger.info(\"Client all train loss: {}\".format(np.mean(clientTrainAllLoss)))\n",
    "            logger.info(\"Client all test accuracy: {}\".format(np.mean(clientTestAllAccuracy)))\n",
    "#             logger.info(\"Client all test loss: {}\".format(np.mean(clientTestAllLoss)))\n",
    "\n",
    "            cleintMacroVal_f1AllTrain.append(np.mean(tempAllMacroVal_f1Train))\n",
    "            clientMacroVal_f1AllTest.append(np.mean(tempAllMacroVal_f1Test))\n",
    "        \n",
    "        # BBP_MAP step\n",
    "        hungarian_weights, assignments_list = BBP_MAP(\n",
    "            retrained_nets, model_meta_data, layer_type, net_dataidx_map, averaging_weights, args, device=device)\n",
    "        \n",
    "        batch_weights = [copy.deepcopy(hungarian_weights)\n",
    "                         for _ in range(args.n_nets)]\n",
    "        \n",
    "        \n",
    "        \n",
    "        serverAccTest, serverLsTest , serverAccTrain , serverLsTrain,sMacroVal_f1Train,sMacroVal_f1Test = compute_full_cnn_accuracy(None,\n",
    "                                      hungarian_weights,\n",
    "                                      train_dl_global,\n",
    "                                      test_dl_global,\n",
    "                                      n_classes,\n",
    "                                      device=device,\n",
    "                                      args=args)\n",
    "        \n",
    "        serverAccuracyTest.append(serverAccTest)\n",
    "        serverLossTest.append(serverLsTest)\n",
    "        serverAccuracyTrain.append(serverAccTrain)\n",
    "        serverLossTrain.append(serverLsTrain)\n",
    "        \n",
    "        serverMacroVal_f1Train.append(sMacroVal_f1Train)\n",
    "        serverMacroVal_f1Test.append(sMacroVal_f1Test)\n",
    "        \n",
    "        \n",
    "        \n",
    "        logger.info(\"Server train accuracy: {}\".format(serverAccTrain))\n",
    "#         logger.info(\"Server train loss: {}\".format(serverLsTrain))\n",
    "        logger.info(\"Server test accuracy: {}\".format(serverAccTest))\n",
    "#         logger.info(\"Server test loss: {}\".format(serverLsTest))\n",
    "                \n",
    "        del hungarian_weights\n",
    "        del retrained_nets\n",
    "        timePerRound = time.time() - startTime \n",
    "        print(timePerRound)\n",
    "        \n",
    "    serverStats = [serverAccuracyTrain,serverLossTrain,serverAccuracyTest,serverLossTest,serverMacroVal_f1Train,serverMacroVal_f1Test]\n",
    "    clientSingleStats = [clientTrainSingleAccuracy,clientTrainSingleAccuracyStd,clientTrainSingleLoss,clientTrainSingleLossStd,clientTestSingleAccuracy,clientTestSingleAccuracyStd,clientTestSingleLoss,clientTestSingleLossStd,cleintMacroVal_f1SingleTrain,clientMacroVal_f1SingleTest]\n",
    "    clientAllStats = [clientTestAllAccuracy,clientTestAllAccuracyStd,clientTestAllLoss,clientTestAllLossStd,clientTrainAllAccuracy,clientTrainAllAccuracyStd,clientTrainAllLoss,clientTrainAllLossStd,cleintMacroVal_f1AllTrain,clientMacroVal_f1AllTest]\n",
    "    return serverStats,clientSingleStats,clientAllStats,batch_weights,timePerRound"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def local_retrain(local_datasets, weights, args, mode=\"bottom-up\", freezing_index=0, ori_assignments=None, device=\"cpu\"):\n",
    "    \"\"\"\n",
    "    freezing_index :: starting from which layer we update the model weights,\n",
    "                      i.e. freezing_index = 0 means we train the whole network normally\n",
    "                           freezing_index = len(model) means we freez the entire network\n",
    "    \"\"\"\n",
    "    input_channel = 6\n",
    "    num_filters = [weights[0].shape[0]]\n",
    "    input_dim = weights[2].shape[0]\n",
    "    hidden_dims = [weights[2].shape[1]]\n",
    "    matched_cnn = SimpleCNNContainer(input_channel = input_channel,num_filters = num_filters,\n",
    "                                       kernel_size=kernelSize,\n",
    "                                       input_dim=input_dim,\n",
    "                                       hidden_dims=hidden_dims,\n",
    "                                       output_dim=activityCount)\n",
    "\n",
    "    new_state_dict = {}\n",
    "    model_counter = 0\n",
    "    n_layers = int(len(weights) / 2)\n",
    "    \n",
    "#     for i in range(np.asarray(weights).shape[0]):\n",
    "#         logger.info(\"weights {}\".format(np.asarray(weights)[i].shape))\n",
    "\n",
    "    # we hardcoded this for now: will probably make changes later\n",
    "    # if mode != \"block-wise\":\n",
    "    if mode not in (\"block-wise\", \"squeezing\"):\n",
    "        __non_loading_indices = []\n",
    "\n",
    "    def __reconstruct_weights(weight, assignment, layer_ori_shape, matched_num_filters=None, weight_type=\"conv_weight\", slice_dim=\"filter\"):\n",
    "        # what contains in the param `assignment` is the assignment for a certain layer, a certain worker\n",
    "        \"\"\"\n",
    "        para:: slice_dim: for reconstructing the conv layers, for each of the three consecutive layers, we need to slice the \n",
    "               filter/kernel to reconstruct the first conv layer; for the third layer in the consecutive block, we need to \n",
    "               slice the\n",
    "               color channel \n",
    "        \"\"\"\n",
    "        if weight_type == \"conv_weight\":\n",
    "            if slice_dim == \"filter\":\n",
    "                res_weight = weight[assignment, :]\n",
    "            elif slice_dim == \"channel\":\n",
    "                _ori_matched_shape = list(copy.deepcopy(layer_ori_shape))\n",
    "                _ori_matched_shape[1] = matched_num_filters\n",
    "                trans_weight = trans_next_conv_layer_forward(\n",
    "                    weight, _ori_matched_shape)\n",
    "                sliced_weight = trans_weight[assignment, :]\n",
    "                res_weight = trans_next_conv_layer_backward(\n",
    "                    sliced_weight, layer_ori_shape)\n",
    "        elif weight_type == \"bias\":\n",
    "            res_weight = weight[assignment]\n",
    "        elif weight_type == \"first_fc_weight\":\n",
    "            # NOTE: please note that in this case, we pass the `estimated_shape` to `layer_ori_shape`:\n",
    "            __ori_shape = weight.shape\n",
    "\n",
    "            res_weight = weight.reshape(\n",
    "                matched_num_filters, layer_ori_shape[2]*__ori_shape[1])[assignment, :]\n",
    "            res_weight = res_weight.reshape(\n",
    "                (len(assignment)*layer_ori_shape[2], __ori_shape[1]))\n",
    "        elif weight_type == \"fc_weight\":\n",
    "            if slice_dim == \"filter\":\n",
    "                res_weight = weight.T[assignment, :]\n",
    "                #res_weight = res_weight.T\n",
    "            elif slice_dim == \"channel\":\n",
    "                res_weight = weight[assignment, :].T\n",
    "        return res_weight  \n",
    "\n",
    "    for param_idx, (key_name, param) in enumerate(matched_cnn.state_dict().items()):\n",
    "        if (param_idx in __non_loading_indices) and (freezing_index[0] != n_layers):\n",
    "            # we need to reconstruct the weights here s.t.\n",
    "            # i) shapes of the weights are euqal to the shapes of the weight in original model (before matching)\n",
    "            # ii) each neuron comes from the corresponding global neuron\n",
    "            _matched_weight = weights[param_idx]\n",
    "            _matched_num_filters = weights[__non_loading_indices[0]].shape[0]\n",
    "            #\n",
    "            # we now use this `_slice_dim` for both conv layers and fc layers\n",
    "            if __non_loading_indices.index(param_idx) != 2:\n",
    "                # please note that for biases, it doesn't really matter if we're going to use filter or channel\n",
    "                _slice_dim = \"filter\"\n",
    "            else:\n",
    "                _slice_dim = \"channel\"\n",
    "            logger.info(\"_slice_dim {}\".format(_slice_dim))\n",
    "#             _slice_dim = \"channel\"\n",
    "\n",
    "            if \"conv\" in key_name or \"features\" in key_name:\n",
    "                if \"weight\" in key_name:\n",
    "                    _res_weight = __reconstruct_weights(weight=_matched_weight, assignment=ori_assignments,\n",
    "                                                        layer_ori_shape=param.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                        weight_type=\"conv_weight\", slice_dim=_slice_dim)\n",
    "                    temp_dict = {key_name: torch.from_numpy(\n",
    "                        _res_weight.reshape(param.size()))}\n",
    "                elif \"bias\" in key_name:\n",
    "                    _res_bias = __reconstruct_weights(weight=_matched_weight, assignment=ori_assignments,\n",
    "                                                      layer_ori_shape=param.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                      weight_type=\"bias\", slice_dim=_slice_dim)\n",
    "                    temp_dict = {key_name: torch.from_numpy(_res_bias)}\n",
    "            elif \"fc\" in key_name or \"classifier\" in key_name:\n",
    "                if \"weight\" in key_name:\n",
    "                    if freezing_index[0] != 4:\n",
    "                        _res_weight = __reconstruct_weights(weight=_matched_weight, assignment=ori_assignments,\n",
    "                                                            layer_ori_shape=param.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                            weight_type=\"fc_weight\", slice_dim=_slice_dim)\n",
    "                        temp_dict = {key_name: torch.from_numpy(_res_weight)}\n",
    "                    else:\n",
    "                        # that's for handling the first fc layer that is connected to the conv blocks\n",
    "                        _res_weight = __reconstruct_weights(weight=_matched_weight, assignment=ori_assignments,\n",
    "                                                            layer_ori_shape=estimated_output.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                            weight_type=\"first_fc_weight\", slice_dim=_slice_dim)\n",
    "                        temp_dict = {key_name: torch.from_numpy(_res_weight.T)}\n",
    "                elif \"bias\" in key_name:\n",
    "                    _res_bias = __reconstruct_weights(weight=_matched_weight, assignment=ori_assignments,\n",
    "                                                      layer_ori_shape=param.size(), matched_num_filters=_matched_num_filters,\n",
    "                                                      weight_type=\"bias\", slice_dim=_slice_dim)\n",
    "                    temp_dict = {key_name: torch.from_numpy(_res_bias)}\n",
    "        else:\n",
    "            if \"conv\" in key_name or \"features\" in key_name:\n",
    "                if \"weight\" in key_name:\n",
    "                    temp_dict = {key_name: torch.from_numpy(\n",
    "                        weights[param_idx].reshape(param.size()))}\n",
    "                elif \"bias\" in key_name:\n",
    "                    temp_dict = {key_name: torch.from_numpy(\n",
    "                        weights[param_idx])}\n",
    "            elif \"fc\" in key_name or \"classifier\" in key_name:\n",
    "                if \"weight\" in key_name:\n",
    "#                     logger.info(\"im in here nara key_name {} weights[param_idx] {}\".format(key_name,np.array(weights[param_idx]).shape))\n",
    "                    temp_dict = {key_name: torch.from_numpy(\n",
    "                        weights[param_idx].T)}\n",
    "                elif \"bias\" in key_name:\n",
    "                    temp_dict = {key_name: torch.from_numpy(\n",
    "                        weights[param_idx])}\n",
    "        new_state_dict.update(temp_dict)\n",
    "        \n",
    "\n",
    "    matched_cnn.load_state_dict(new_state_dict)\n",
    "\n",
    "                    \n",
    "                \n",
    "    for param_idx, param in enumerate(matched_cnn.parameters()):\n",
    "        if mode == \"bottom-up\":\n",
    "            # for this freezing mode, we freeze the layer before freezing index\n",
    "            if param_idx < freezing_index:\n",
    "                param.requires_grad = False\n",
    "        elif mode == \"per-layer\":\n",
    "            # for this freezing mode, we only unfreeze the freezing index\n",
    "            if param_idx not in (2*freezing_index-2, 2*freezing_index-1):\n",
    "                param.requires_grad = False\n",
    "        elif mode == \"block-wise\":\n",
    "            # for block-wise retraining the `freezing_index` becomes a range of indices\n",
    "            if param_idx not in __non_loading_indices:\n",
    "                param.requires_grad = False\n",
    "        elif mode == \"squeezing\":\n",
    "            pass\n",
    "\n",
    "    matched_cnn.to(device).train()\n",
    "    # start training last fc layers:\n",
    "    train_dl_local = local_datasets[0]\n",
    "    test_dl_local = local_datasets[1]\n",
    "\n",
    "    if mode != \"block-wise\":\n",
    "        if freezing_index < (len(weights) - 2):\n",
    "            optimizer_fine_tune = optim.SGD(filter(\n",
    "                lambda p: p.requires_grad, matched_cnn.parameters()), lr=args.retrain_lr, momentum=0.9)\n",
    "        else:\n",
    "            optimizer_fine_tune = optim.SGD(filter(lambda p: p.requires_grad, matched_cnn.parameters(\n",
    "            )), lr=(args.retrain_lr/10), momentum=0.9, weight_decay=0.0001)\n",
    "    else:\n",
    "        #optimizer_fine_tune = optim.SGD(filter(lambda p: p.requires_grad, matched_cnn.parameters()), lr=args.retrain_lr, momentum=0.9)\n",
    "        optimizer_fine_tune = optim.Adam(filter(lambda p: p.requires_grad, matched_cnn.parameters(\n",
    "        )), lr=0.001, weight_decay=0.0001, amsgrad=True)\n",
    "\n",
    "    criterion_fine_tune = nn.CrossEntropyLoss().to(device)\n",
    "\n",
    "\n",
    "    if mode != \"block-wise\":\n",
    "        if freezing_index < (len(weights) - 2):\n",
    "            retrain_epochs = args.retrain_epochs\n",
    "        else:\n",
    "            retrain_epochs = int(args.retrain_epochs*3)\n",
    "    else:\n",
    "        retrain_epochs = args.retrain_epochs\n",
    "\n",
    "    for epoch in range(retrain_epochs):\n",
    "        epoch_loss_collector = []\n",
    "        for batch_idx, (x, target) in enumerate(train_dl_local):\n",
    "            x, target = x.to(device), target.to(device)\n",
    "\n",
    "            optimizer_fine_tune.zero_grad()\n",
    "            x.requires_grad = True\n",
    "            target.requires_grad = False\n",
    "            target = target.long()\n",
    "\n",
    "            out = matched_cnn(x)\n",
    "            loss = criterion_fine_tune(out, target)\n",
    "            epoch_loss_collector.append(loss.item())\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer_fine_tune.step()\n",
    "\n",
    "        epoch_loss = sum(epoch_loss_collector) / len(epoch_loss_collector)\n",
    "\n",
    "\n",
    "    return matched_cnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model(model, model_index):\n",
    "    logger.info(\"saving local model-{}\".format(model_index))\n",
    "    with open(\"trained_local_model\"+str(model_index), \"wb\") as f_:\n",
    "        torch.save(model.state_dict(), f_)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_uniform_weights(n_classes, net_cnt, fill_val=1):\n",
    "    weights_list = {}\n",
    "\n",
    "    for net_i in range(net_cnt):\n",
    "        temp = np.array([fill_val] * n_classes, dtype=np.float32)\n",
    "        weights_list[net_i] = torch.from_numpy(temp).view(1, -1)\n",
    "\n",
    "    return weights_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdm_prepare_freq(cls_freqs, n_classes):\n",
    "    freqs = []\n",
    "\n",
    "    for net_i in sorted(cls_freqs.keys()):\n",
    "        net_freqs = [0] * n_classes\n",
    "\n",
    "        for cls_i in cls_freqs[net_i]:\n",
    "            net_freqs[cls_i] = cls_freqs[net_i][cls_i]\n",
    "\n",
    "        freqs.append(np.array(net_freqs))\n",
    "\n",
    "    return freqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trans_next_conv_layer_forward(layer_weight, next_layer_shape):\n",
    "#     logger.info(\"trans_next_conv_layer_forward next_layer_shape {}\".format(next_layer_shape))\n",
    "    reshaped = layer_weight.reshape(next_layer_shape).transpose(\n",
    "        (1, 0, 2)).reshape((next_layer_shape[1], -1))\n",
    "    return reshaped\n",
    "\n",
    "\n",
    "def trans_next_conv_layer_backward(layer_weight, next_layer_shape):\n",
    "#     logger.info(\"trans_next_conv_layer_backward next_layer_shape {}\".format(next_layer_shape))\n",
    "    reconstructed_next_layer_shape = (\n",
    "        next_layer_shape[1], next_layer_shape[0], next_layer_shape[2])\n",
    "    reshaped = layer_weight.reshape(reconstructed_next_layer_shape).transpose(\n",
    "        1, 0, 2).reshape(next_layer_shape[0], -1)\n",
    "    return reshaped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(global_weights, global_sigmas):\n",
    "    obj = ((global_weights)/ global_sigmas).sum()\n",
    "    return obj\n",
    "\n",
    "\n",
    "def patch_weights(w_j, L_next, assignment_j_c):\n",
    "    if assignment_j_c is None:\n",
    "        return w_j\n",
    "    new_w_j = np.zeros((w_j.shape[0], L_next))\n",
    "    new_w_j[:, assignment_j_c] = w_j\n",
    "    return new_w_j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args.model = \"simple-cnn\"\n",
    "logger.info(\"Initializing nets\")\n",
    "nets, model_meta_data, layer_type = init_models(args.n_nets, args)\n",
    "logger.info(\"Retrain? : {}\".format(args.retrain))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# local training stage\n",
    "nets_list = local_train(nets, args, net_dataidx_map, device=device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl_global, test_dl_global = get_dataloader(\n",
    "    args.dataset, args_datadir, args.batch_size, 32)\n",
    "\n",
    "# ensemble part of experiments\n",
    "logger.info(\"Computing Uniform ensemble accuracy\")\n",
    "uens_train_acc, _ = compute_ensemble_accuracy(\n",
    "    nets_list, train_dl_global, n_classes,  uniform_weights=True, device=device)\n",
    "uens_test_acc, _ = compute_ensemble_accuracy(\n",
    "    nets_list, test_dl_global, n_classes, uniform_weights=True, device=device)\n",
    "\n",
    "logger.info(\"Uniform ensemble (Train acc): {}\".format(uens_train_acc))\n",
    "logger.info(\"Uniform ensemble (Test acc): {}\".format(uens_test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_weights = pdm_prepare_full_weights_cnn(nets_list, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hungarian_weights, assignments_list = BBP_MAP(\n",
    "    nets_list, model_meta_data, layer_type, net_dataidx_map, averaging_weights, args, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.info(\"Weights shapes: {}\".format(\n",
    "    [bw.shape for bw in hungarian_weights]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_weights = pdm_prepare_full_weights_cnn(nets_list, device=device)\n",
    "total_data_points = sum([len(net_dataidx_map[r])\n",
    "                         for r in range(args.n_nets)])\n",
    "fed_avg_freqs = [len(net_dataidx_map[r]) /\n",
    "                 total_data_points for r in range(args.n_nets)]\n",
    "logger.info(\"Total data points: {}\".format(total_data_points))\n",
    "logger.info(\"Freq of FedAvg: {}\".format(fed_avg_freqs))\n",
    "\n",
    "averaged_weights = []\n",
    "num_layers = len(batch_weights[0])\n",
    "for i in range(num_layers):\n",
    "    avegerated_weight = sum([b[i] * fed_avg_freqs[j]\n",
    "                             for j, b in enumerate(batch_weights)])\n",
    "    averaged_weights.append(avegerated_weight)\n",
    "\n",
    "for aw in averaged_weights:\n",
    "    logger.info(aw.shape)\n",
    "\n",
    "models = nets_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_model_averaging_accuracy(models,\n",
    "                                     averaged_weights,\n",
    "                                     train_dl_global,\n",
    "                                     test_dl_global,\n",
    "                                     n_classes,\n",
    "                                     args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_full_cnn_accuracy(models, weights, train_dl, test_dl, n_classes, args, device=\"cpu\",clientServer=\"Client\"):\n",
    "    \"\"\"Note that we only handle the FC weights for now\"\"\"\n",
    "    # we need to figure out the FC dims first\n",
    "\n",
    "    input_channel = 6\n",
    "\n",
    "    num_filters = [weights[0].shape[0]]\n",
    "    input_dim = weights[2].shape[0]\n",
    "    hidden_dims = [weights[2].shape[1]]\n",
    "    matched_cnn = SimpleCNNContainer(input_channel=input_channel,\n",
    "                                     num_filters=num_filters,\n",
    "                                     kernel_size=kernelSize,\n",
    "                                     input_dim=input_dim,\n",
    "                                     hidden_dims=hidden_dims,\n",
    "                                     output_dim=activityCount)\n",
    "    \n",
    "#     for index, wei in enumerate(weights):\n",
    "#         logger.info(\"weight shapes {}\".format(np.asarray(wei).shape))\n",
    "    if(clientServer == \"test\"):\n",
    "        matched_cnn = models\n",
    "    #logger.info(\"Keys of layers of convblock ...\")\n",
    "    new_state_dict = {}\n",
    "    model_counter = 0\n",
    "    # handle the conv layers part which is not changing\n",
    "    for param_idx, (key_name, param) in enumerate(matched_cnn.state_dict().items()):\n",
    "        # print(\"&\"*30)\n",
    "        #print(\"Key: {}, Weight Shape: {}, Matched weight shape: {}\".format(key_name, param.size(), weights[param_idx].shape))\n",
    "        # print(\"&\"*30)\n",
    "        if \"conv\" in key_name or \"features\" in key_name:\n",
    "            if \"weight\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(\n",
    "                    weights[param_idx].reshape(param.size()))}\n",
    "            elif \"bias\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(weights[param_idx])}\n",
    "        elif \"fc\" in key_name or \"classifier\" in key_name:\n",
    "            if \"weight\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(weights[param_idx].T)}\n",
    "            elif \"bias\" in key_name:\n",
    "                temp_dict = {key_name: torch.from_numpy(weights[param_idx])}\n",
    "\n",
    "        new_state_dict.update(temp_dict)\n",
    "    matched_cnn.load_state_dict(new_state_dict)\n",
    "    matched_cnn.to(device)\n",
    "    matched_cnn.eval()\n",
    "\n",
    "    ##\n",
    "    criterion = nn.CrossEntropyLoss().to(device)\n",
    "    correct, total = 0, 0\n",
    "    total_loss = []\n",
    "    allPred = []\n",
    "    for batch_idx, (x, target) in enumerate(test_dl):\n",
    "        x, target = x.to(device), target.to(device)\n",
    "        out_k = matched_cnn(x)\n",
    "        _, pred_label = torch.max(out_k, 1)\n",
    "        total += x.data.size()[0]\n",
    "        correct += (pred_label == target.data).sum().item()\n",
    "#         pred_label = pred_label.squeeze_()\n",
    "        allPred.append(pred_label)\n",
    "        loss = criterion(out_k, target.long())\n",
    "        total_loss.append(loss.item())\n",
    "        \n",
    "    allPred = torch.cat(allPred)\n",
    "\n",
    "    macroVal_f1Test = f1_score(test_dl.dataset.target,allPred.cpu(),average='macro')\n",
    "\n",
    "        \n",
    "    clientAccuracyTest = correct/total\n",
    "    clientLossTest = sum(total_loss) / len(total_loss)\n",
    "    \n",
    "    allPred = []\n",
    "    correct, total = 0, 0\n",
    "    total_loss = []\n",
    "    for batch_idx, (x, target) in enumerate(train_dl):\n",
    "        x, target = x.to(device), target.to(device)\n",
    "        out_k = matched_cnn(x)\n",
    "        _, pred_label = torch.max(out_k, 1)\n",
    "        total += x.data.size()[0]\n",
    "        correct += (pred_label == target.data).sum().item()\n",
    "#         pred_label = pred_label.squeeze_()\n",
    "        allPred.append(pred_label)\n",
    "        loss = criterion(out_k, target.long())\n",
    "        total_loss.append(loss.item())\n",
    "                \n",
    "    allPred = torch.cat(allPred)\n",
    "    macroVal_f1Train = f1_score(train_dl.dataset.target,allPred.cpu(),average='macro')\n",
    "\n",
    "\n",
    "    clientAccuracyTrain = correct/total\n",
    "\n",
    "\n",
    "    clientLossTrain = sum(total_loss) / len(total_loss)    \n",
    "    \n",
    "    return clientAccuracyTest,clientLossTest,clientAccuracyTrain,clientLossTrain,macroVal_f1Train,macroVal_f1Test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comm_init_batch_weights = [copy.deepcopy(\n",
    "    hungarian_weights) for _ in range(args.n_nets)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meanHistoryDist = []\n",
    "stdHistoryDist = []\n",
    "\n",
    "meanRoundLayerHistory = []\n",
    "stdRoundLayerHistory = []\n",
    "\n",
    "meanRoundGeneralLayerHistory = []\n",
    "stdRoundGeneralLayerHistory = []\n",
    "\n",
    "roundTime = 0\n",
    "\n",
    "serverStats,clientSingleStats,clientAllStats,finalWeights,roundTime = fedma_comm(comm_init_batch_weights,\n",
    "           model_meta_data, layer_type, net_dataidx_map,\n",
    "           averaging_weights, args,\n",
    "           train_dl_global,\n",
    "           test_dl_global,\n",
    "           assignments_list,\n",
    "           comm_round=args.comm_round,\n",
    "           device=device)\n",
    "\n",
    "\n",
    "serverAccuracy = serverStats[2]\n",
    "serverLoss = serverStats[3]\n",
    "serverAccuracyTrain = serverStats[0]\n",
    "serverLossTrain = serverStats[1]\n",
    "\n",
    "serverMacroTrain = serverStats[4]\n",
    "serverMacroTest = serverStats[5]\n",
    "\n",
    "clientAccuracy = clientSingleStats[4]\n",
    "clientAccuracyStd = clientSingleStats[5]\n",
    "clientLoss = clientSingleStats[6]\n",
    "clientLossStd = clientSingleStats[7]\n",
    "clientAccuracyTrain = clientSingleStats[0]\n",
    "clientAccuracyStdTrain = clientSingleStats[1]\n",
    "clientLossTrain = clientSingleStats[2]\n",
    "clientLossStdTrain = clientSingleStats[3]\n",
    "\n",
    "clientMacroTrain = clientSingleStats[8]\n",
    "clientMacroTest = clientSingleStats[9]\n",
    "\n",
    "\n",
    "clientAllAccuracy = clientAllStats[4]\n",
    "clientAllAccuracyStd = clientAllStats[5]\n",
    "clientAllLoss = clientAllStats[6]\n",
    "clientAllLossStd = clientAllStats[7]\n",
    "clientAllAccuracyTrain = clientAllStats[0]\n",
    "clientAllAccuracyStdTrain = clientAllStats[1]\n",
    "clientAllLossTrain = clientAllStats[2]\n",
    "clientAllLossStdTrain = clientAllStats[3]\n",
    "\n",
    "clientAllMacroTrain = clientAllStats[8]\n",
    "clientAllMacroTest = clientAllStats[9]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serverAccuracyTrain = np.asarray(serverAccuracyTrain)\n",
    "serverLossTrain = np.asarray(serverLossTrain)\n",
    "serverAccuracy = np.asarray(serverAccuracy)\n",
    "serverLoss = np.asarray(serverLoss)\n",
    "\n",
    "clientAccuracy = np.asarray(clientAccuracy)\n",
    "clientAccuracyStd = np.asarray(clientAccuracyStd)\n",
    "clientLoss = np.asarray(clientLoss)\n",
    "clientLossStd = np.asarray(clientLossStd)\n",
    "clientAccuracyTrain = np.asarray(clientAccuracyTrain)\n",
    "clientAccuracyStdTrain = np.asarray(clientAccuracyStdTrain)\n",
    "clientLossTrain = np.asarray(clientLossTrain)\n",
    "clientLossStdTrain = np.asarray(clientLossStdTrain)\n",
    "\n",
    "clientAllAccuracy = np.asarray(clientAllAccuracy)\n",
    "clientAllAccuracyStd = np.asarray(clientAllAccuracyStd)\n",
    "clientAllLoss = np.asarray(clientAllLoss)\n",
    "clientAllLossStd = np.asarray(clientAllLossStd)\n",
    "clientAllAccuracyTrain = np.asarray(clientAllAccuracyTrain)\n",
    "clientAllAccuracyStdTrain = np.asarray(clientAllAccuracyStdTrain)\n",
    "clientAllLossTrain = np.asarray(clientAllLossTrain)\n",
    "clientAllLossStdTrain = np.asarray(clientAllLossStdTrain)\n",
    "\n",
    "\n",
    "\n",
    "clientAllMacroTrain = np.asarray(clientAllMacroTrain)\n",
    "clientAllMacroTest = np.asarray(clientAllMacroTest)\n",
    "\n",
    "clientMacroTrain = np.asarray(clientMacroTrain)\n",
    "clientMacroTest = np.asarray(clientMacroTest)\n",
    "\n",
    "serverMacroTrain = np.asarray(serverMacroTrain)\n",
    "serverMacroTest = np.asarray(serverMacroTest)\n",
    "\n",
    "if(euclid):\n",
    "    meanHistoryDist = np.asarray(meanHistoryDist).T\n",
    "    stdHistoryDist = np.asarray(stdHistoryDist).T\n",
    "    meanRoundLayerHistory = np.asarray(meanRoundLayerHistory).T\n",
    "    stdRoundLayerHistory = np.asarray(stdRoundLayerHistory).T\n",
    "    meanRoundGeneralLayerHistory = np.asarray(meanRoundGeneralLayerHistory)\n",
    "    stdRoundGeneralLayerHistory = np.asarray(stdRoundGeneralLayerHistory)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(filepath+\"trainingStats\", exist_ok=True)\n",
    "hkl.dump(serverAccuracy,filepath + \"trainingStats/serverAccuracy.hkl\" )\n",
    "hkl.dump(serverLoss,filepath + \"trainingStats/serverLoss.hkl\" )\n",
    "\n",
    "hkl.dump(serverAccuracyTrain,filepath + \"trainingStats/serverAccuracyTrain.hkl\" )\n",
    "hkl.dump(serverLossTrain,filepath + \"trainingStats/serverLossTrain.hkl\" )\n",
    "\n",
    "hkl.dump(clientAccuracy,filepath + \"trainingStats/clientAccuracy.hkl\" )\n",
    "hkl.dump(clientAccuracyStd,filepath + \"trainingStats/clientAccuracyStd.hkl\" )\n",
    "\n",
    "hkl.dump(clientLoss,filepath + \"trainingStats/clientLoss.hkl\" )\n",
    "hkl.dump(clientLossStd,filepath + \"trainingStats/clientLossStd.hkl\" )\n",
    "\n",
    "hkl.dump(clientAccuracyTrain,filepath + \"trainingStats/clientAccuracyTrain.hkl\" )\n",
    "hkl.dump(clientAccuracyStdTrain,filepath + \"trainingStats/clientAccuracyStdTrain.hkl\" )\n",
    "\n",
    "hkl.dump(clientLossTrain,filepath + \"trainingStats/clientLossTrain.hkl\" )\n",
    "hkl.dump(clientLossStdTrain,filepath + \"trainingStats/clientLossStdTrain.hkl\" )\n",
    "\n",
    "hkl.dump(clientAllAccuracy,filepath + \"trainingStats/clientAllAccuracy.hkl\" )\n",
    "hkl.dump(clientAllAccuracyStd,filepath + \"trainingStats/clientAllAccuracyStd.hkl\" )\n",
    "\n",
    "hkl.dump(clientAllLoss,filepath + \"trainingStats/clientAllLoss.hkl\" )\n",
    "hkl.dump(clientAllLossStd,filepath + \"trainingStats/clientAllLossStd.hkl\" )\n",
    "\n",
    "hkl.dump(clientAllAccuracyTrain,filepath + \"trainingStats/clientAllAccuracyTrain.hkl\" )\n",
    "hkl.dump(clientAllAccuracyStdTrain,filepath + \"trainingStats/clientAllAccuracyStdTrain.hkl\" )\n",
    "\n",
    "hkl.dump(clientAllLossTrain,filepath + \"trainingStats/clientAllLossTrain.hkl\" )\n",
    "hkl.dump(clientAllLossStdTrain,filepath + \"trainingStats/clientAllLossStdTrain.hkl\" )\n",
    "\n",
    "hkl.dump(clientAllMacroTrain,filepath + \"trainingStats/clientAllMacroTrain.hkl\" )\n",
    "hkl.dump(clientAllMacroTest,filepath + \"trainingStats/clientAllMacroTest.hkl\" )\n",
    "\n",
    "hkl.dump(clientMacroTrain,filepath + \"trainingStats/clientMacroTrain.hkl\" )\n",
    "hkl.dump(clientMacroTest,filepath + \"trainingStats/clientMacroTest.hkl\" )\n",
    "\n",
    "hkl.dump(serverMacroTrain,filepath + \"trainingStats/serverMacroTrain.hkl\" )\n",
    "hkl.dump(serverMacroTest,filepath + \"trainingStats/serverMacroTest.hkl\" )\n",
    "\n",
    "if(euclid):\n",
    "    hkl.dump(meanHistoryDist,filepath + \"trainingStats/meanHistoryDist.hkl\" )\n",
    "    hkl.dump(stdHistoryDist,filepath + \"trainingStats/stdHistoryDist.hkl\" )\n",
    "    hkl.dump(meanRoundLayerHistory,filepath + \"trainingStats/meanRoundLayerHistory.hkl\" )\n",
    "    hkl.dump(stdRoundLayerHistory,filepath + \"trainingStats/stdRoundLayerHistory.hkl\" )\n",
    "    hkl.dump(meanRoundGeneralLayerHistory,filepath + \"trainingStats/meanRoundGeneralLayerHistory.hkl\" )\n",
    "    hkl.dump(stdRoundGeneralLayerHistory,filepath + \"trainingStats/stdRoundGeneralLayerHistory.hkl\" )\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch_range = range(1, args.comm_round+1)\n",
    "\n",
    "if(algorithm != \"FEDPER\"):\n",
    "    plt.plot(epoch_range, serverAccuracyTrain, label = 'Server Train')\n",
    "    plt.plot(epoch_range, serverAccuracy, label= 'Server Test')\n",
    "   \n",
    "plt.errorbar(epoch_range, clientAccuracyTrain, yerr=clientAccuracyStdTrain, label='Client Own Train', alpha = 0.6)\n",
    "plt.errorbar(epoch_range, clientAccuracy, yerr=clientAccuracyStd, label='Client Own Test',alpha = 0.6)\n",
    "\n",
    "if(ClientAllTest  == True):\n",
    "    plt.errorbar(epoch_range, clientAllAccuracyTrain, yerr=clientAllAccuracyStdTrain, label='Client All Train', alpha = 0.6)\n",
    "    plt.errorbar(epoch_range, clientAllAccuracy, yerr=clientAllAccuracyStd, label='Client All Test', alpha = 0.6)\n",
    "    \n",
    "    plt.plot(epoch_range, clientAllAccuracyTrain,markevery=[np.argmax(clientAllAccuracyTrain)], ls=\"\", marker=\"o\",color=\"purple\")\n",
    "    plt.plot(epoch_range, clientAllAccuracy,markevery=[np.argmax(clientAllAccuracy)], ls=\"\", marker=\"o\",color=\"brown\")  \n",
    "\n",
    "plt.plot(epoch_range, serverAccuracyTrain,markevery=[np.argmax(serverAccuracyTrain)], ls=\"\", marker=\"o\",color=\"blue\")\n",
    "plt.plot(epoch_range, serverAccuracy,markevery=[np.argmax(serverAccuracy)], ls=\"\", marker=\"o\",color=\"orange\") \n",
    "\n",
    "plt.plot(epoch_range, clientAccuracyTrain,markevery=[np.argmax(clientAccuracyTrain)], ls=\"\", marker=\"o\",color=\"green\")\n",
    "plt.plot(epoch_range, clientAccuracy,markevery=[np.argmax(clientAccuracy)], ls=\"\", marker=\"o\",color=\"red\")  \n",
    "    \n",
    "plt.title('Model accuracy')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.xlabel('Communication Round')\n",
    "plt.legend(loc='lower right')\n",
    "plt.savefig(filepath+'LearningAccuracy.png', dpi=100)\n",
    "plt.clf()\n",
    "\n",
    "if(algorithm != \"FEDPER\"):\n",
    "    plt.plot(epoch_range, serverLossTrain, label = 'Server Train')\n",
    "    plt.plot(epoch_range, serverLoss, label= 'Server Test')\n",
    "\n",
    "plt.errorbar(epoch_range, clientLossTrain, yerr=clientLossStdTrain, label='Client Own Train',alpha = 0.6)\n",
    "plt.errorbar(epoch_range, clientLoss, yerr=clientLossStd, label='Client Own Test',alpha = 0.6)\n",
    "    \n",
    "if(ClientAllTest  == True):\n",
    "    plt.errorbar(epoch_range, clientAllLossTrain, yerr=clientAllLossStdTrain, label='Client All Train',alpha = 0.6)\n",
    "    plt.errorbar(epoch_range, clientAllLoss, yerr=clientAllLossStd, label='Client All Test',alpha = 0.6)\n",
    "    plt.plot(epoch_range, clientAllLossTrain,markevery=[np.argmin(clientAllLossTrain)], ls=\"\", marker=\"o\",color=\"purple\")\n",
    "    plt.plot(epoch_range, clientAllLoss,markevery=[np.argmin(clientAllLoss)], ls=\"\", marker=\"o\",color=\"brown\")  \n",
    "    \n",
    "plt.plot(epoch_range, serverLossTrain,markevery=[np.argmin(serverLossTrain)], ls=\"\", marker=\"o\",color=\"blue\")\n",
    "plt.plot(epoch_range, serverLoss,markevery=[np.argmin(serverLoss)], ls=\"\", marker=\"o\",color=\"orange\") \n",
    "\n",
    "plt.plot(epoch_range, clientLossTrain,markevery=[np.argmin(clientLossTrain)], ls=\"\", marker=\"o\",color=\"green\")\n",
    "plt.plot(epoch_range, clientLoss,markevery=[np.argmin(clientLoss)], ls=\"\", marker=\"o\",color=\"red\")  \n",
    "\n",
    "plt.title('Model Loss')\n",
    "plt.ylabel('Loss')\n",
    "plt.xlabel('Communication Round')\n",
    "plt.legend(loc='upper right')\n",
    "\n",
    "plt.savefig(filepath+'LearningLoss.png', dpi=100)\n",
    "plt.clf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if(euclid):\n",
    "    for i in range(clientCount):\n",
    "        plt.errorbar(epoch_range, meanHistoryDist[i], yerr=stdHistoryDist[i], label='Client '+str(i+1))\n",
    "    plt.title('Distance between client & server model')\n",
    "    plt.ylabel('Euclidiance Distance')\n",
    "    plt.xlabel('Communication Round')\n",
    "    plt.savefig(filepath+'allClientEuclid.png', dpi=100)\n",
    "#     plt.legend(loc='upper right')\n",
    "    plt.clf()\n",
    "    \n",
    "    for i in range(len(layerType)):\n",
    "        plt.errorbar(epoch_range, meanRoundLayerHistory[i], yerr=stdRoundLayerHistory[i], label='Layer '+str(i+1)) \n",
    "    plt.errorbar(epoch_range, meanRoundGeneralLayerHistory, yerr=stdRoundGeneralLayerHistory, label='General')\n",
    "    plt.title('Layer distance between client & server model')\n",
    "    plt.ylabel('Euclidiance Distance')\n",
    "    plt.xlabel('Communication Round')\n",
    "    plt.savefig(filepath+'LayerClientEuclid.png', dpi=100)\n",
    "    plt.legend(loc='upper right')\n",
    "    plt.clf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hkl.dump(finalWeights,filepath + \"trainingStats/finalWeights.hkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.info(\"Best Server accuracy at {} at CR {}\".format(serverAccuracy.max(),serverAccuracy.argmax()+1))\n",
    "logger.info(\"Best client client accuracy at {} at CR {}\".format(clientAccuracy.max(),clientAccuracy.argmax()+1))\n",
    "logger.info(\"Best all client accuracy at {} at CR {}\".format(clientAllAccuracy.max(),clientAllAccuracy.argmax()+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelStatistics = {\n",
    "    \"Server accuracy\" : roundNumber(serverAccuracy.max()),\n",
    "    \"Best server round:\": serverAccuracy.argmax()+1,\n",
    "    \"Single client accuracy:\" : roundNumber(clientAccuracy.max()),\n",
    "    \"Best single client Round:\" : clientAccuracy.argmax()+1,\n",
    "    \"All client accuracy:\": roundNumber(clientAllAccuracy.max()),\n",
    "    \"Best all client Round\": clientAllAccuracy.argmax()+1,\n",
    "    \"All client F-Measure:\":roundNumber(clientAllMacroTest),\n",
    "    \"Single client F-Measure:\":roundNumber(clientMacroTest),\n",
    "    \"Server F-Measure:\":roundNumber(serverMacroTest),\n",
    "    \"Time per round:\":roundTime\n",
    "}    \n",
    "\n",
    "\n",
    "with open(filepath +'bestAccuracyStats.csv','w') as f:\n",
    "    w = csv.writer(f)\n",
    "    w.writerows(modelStatistics.items())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "modelFinalShape = {}\n",
    "for i in range (np.int(np.floor(np.asarray(finalWeights[0]).shape[0]/2))):\n",
    "    modelFinalShape[\"Layer \"+str(i+1)] = np.asarray(finalWeights[0][i*2+1]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(filepath +'modelFinalShape.csv','w') as f:\n",
    "    w = csv.writer(f)\n",
    "    w.writerows(modelFinalShape.items())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}