fedAVG.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# implementation of FedAVG\n",
    "# Communication-Efficient Learning of Deep Networks from Decentralized Data by H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, Blaise Agüera y Arcas\n",
    "# https://arxiv.org/abs/1602.05629\n",
    "# https://github.com/AshwinRJ/Federated-Learning-PyTorch\n",
    "\n",
    "import copy\n",
    "\n",
    "import torch\n",
    "from torchsummary import summary\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "from models import ResNet50\n",
    "# from utils import get_dataset, average_weights, exp_details\n",
    "# from utils_v2 import get_dataset, average_weights, exp_details\n",
    "from update import LocalUpdate, test_inference"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# parameters\n",
    "iid = 0 # if the data is i.i.d or not\n",
    "unbalanced = 1 # in non i.i.d. setting split the data between clients equally or not\n",
    "num_users = 100 # number of client\n",
    "frac = 0.1 # fraction of the clients to be used for federated updates\n",
    "n_epochs = 30\n",
    "gpu = 0\n",
    "optimizer = \"sgd\" #sgd or adam\n",
    "local_batch_size = 10 # batch size of local updates in each user\n",
    "lr = 0.01 # learning rate\n",
    "local_epochs = 10\n",
    "loss_function = \"CrossEntropyLoss\"\n",
    "\n",
    "percentage = 90  # percentage of strugglers\n",
    "\n",
    "num_groups = 0  # 0 for BatchNorm, > 0 for GroupNorm\n",
    "if num_groups == 0:\n",
    "    normalization_type = \"BatchNorm\"\n",
    "else:\n",
    "    normalization_type = \"GroupNorm\""
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if iid:\n",
    "    from utils_v2 import get_dataset, average_weights, weighted_average_weights, exp_details\n",
    "else:\n",
    "    from utils import get_dataset, average_weights, weighted_average_weights, exp_details"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "exp_details(\"ResNet50\", optimizer, lr, normalization_type, n_epochs, iid, frac,\n",
    "            local_batch_size, local_epochs, unbalanced, num_users)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# for REPRODUCIBILITY https://pytorch.org/docs/stable/notes/randomness.html\n",
    "torch.manual_seed(0)\n",
    "\n",
    "g = torch.Generator()\n",
    "g.manual_seed(0)\n",
    "\n",
    "np.random.seed(0)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "train_dataset, test_dataset, user_groups = get_dataset(iid=iid, unbalanced=unbalanced,\n",
    "                                                       num_users=num_users)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model = ResNet50(n_type=normalization_type)\n",
    "# model = CNNCifar()\n",
    "\n",
    "train_on_gpu = torch.cuda.is_available()\n",
    "\n",
    "if not train_on_gpu:\n",
    "    print('CUDA is not available.  Training on CPU ...')\n",
    "    device = torch.device(\"cpu\")\n",
    "    gpu = 0\n",
    "else:\n",
    "    print('CUDA is available!  Training on GPU ...')\n",
    "    device = torch.device(\"cuda\")\n",
    "    gpu = 1\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "# set the model to train\n",
    "model.train()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# total number of params 591,322\n",
    "summary(model, (3, 32, 32))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def GenerateLocalEpochs(percentage, size, max_epochs):\n",
    "    ''' Method generates list of epochs for selected clients\n",
    "    to replicate system heteroggeneity\n",
    "\n",
    "    Params:\n",
    "      percentage: percentage of clients to have fewer than E epochs\n",
    "      size:       total size of the list\n",
    "      max_epochs: maximum value for local epochs\n",
    "\n",
    "    Returns:\n",
    "      List of size epochs for each Client Update\n",
    "\n",
    "    '''\n",
    "\n",
    "    # if percentage is 0 then each client runs for E epochs\n",
    "    if percentage == 0:\n",
    "        return np.array([max_epochs]*size)\n",
    "    else:\n",
    "        # get the number of clients to have fewer than E epochs\n",
    "        heterogenous_size = int((percentage/100) * size)\n",
    "\n",
    "        # generate random uniform epochs of heterogenous size between 1 and E\n",
    "        epoch_list = np.random.randint(1, max_epochs, heterogenous_size)\n",
    "\n",
    "        # the rest of the clients will have E epochs\n",
    "        remaining_size = size - heterogenous_size\n",
    "        rem_list = [max_epochs]*remaining_size\n",
    "\n",
    "        epoch_list = np.append(epoch_list, rem_list, axis=0)\n",
    "\n",
    "        # shuffle the list and return\n",
    "        np.random.shuffle(epoch_list)\n",
    "\n",
    "        return epoch_list"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# copy weights\n",
    "global_weights = model.state_dict()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "test_acc_list = []  # final accuracy is the mean of the last 10\n",
    "test_loss_list = []\n",
    "\n",
    "# training\n",
    "train_loss, train_accuracy = [], []\n",
    "val_acc_list, net_list = [], []\n",
    "cv_loss, cv_acc = [], []\n",
    "\n",
    "for epoch in range(1, n_epochs+1):\n",
    "    local_weights = []\n",
    "    local_losses = []\n",
    "\n",
    "    print(f'Epoch: {epoch} \\n')\n",
    "\n",
    "\n",
    "    ###################\n",
    "    # train the model #\n",
    "    ###################\n",
    "    model.train()\n",
    "\n",
    "    # different clients at each epoch\n",
    "    m = max(int(frac * num_users), 1) # number of users to be used for federated updates, at least 1\n",
    "    idxs_users = np.random.choice(range(num_users), m, replace=False) # choose randomly m users\n",
    "\n",
    "    # define how many local epochs for each client\n",
    "    heterogenous_epoch_list = GenerateLocalEpochs(percentage, size=m, max_epochs=local_epochs)\n",
    "\n",
    "    stragglers_indices = np.argwhere(heterogenous_epoch_list < local_epochs)  # find strugglers\n",
    "    heterogenous_epoch_list = np.delete(heterogenous_epoch_list, stragglers_indices)  # ignore strugglers\n",
    "    idxs_users = np.delete(idxs_users, stragglers_indices)  # ignore strugglers\n",
    "\n",
    "    for idx, i in zip(idxs_users, range(len(heterogenous_epoch_list))):  # for each user\n",
    "        # get local model\n",
    "        local_model = LocalUpdate(dataset=train_dataset, idxs=user_groups[idx],\n",
    "                                  gpu=gpu, optimizer=optimizer,\n",
    "                                  local_batch_size=local_batch_size, lr=lr,\n",
    "                                  local_epochs=heterogenous_epoch_list[i], loss_function=loss_function)\n",
    "\n",
    "        # get updated weight and loss from local model\n",
    "        w, loss = local_model.update_weights(model=copy.deepcopy(model), # pass the global model to the clients\n",
    "                                             global_round=epoch)\n",
    "\n",
    "        print('| Client : {} | Average Loss: {:.4f} '.format(\n",
    "            idx, loss))\n",
    "\n",
    "        local_weights.append(copy.deepcopy(w))\n",
    "        local_losses.append(copy.deepcopy(loss))\n",
    "\n",
    "    # compute global weights (average of local weights)\n",
    "    if unbalanced:\n",
    "        global_weights = weighted_average_weights(local_weights, user_groups, idxs_users)\n",
    "    else:\n",
    "        global_weights = average_weights(local_weights)\n",
    "\n",
    "    # update weights of the global model\n",
    "    model.load_state_dict(global_weights)\n",
    "\n",
    "    # compute average loss\n",
    "    loss_avg = sum(local_losses) / len(local_losses)\n",
    "    train_loss.append(loss_avg)\n",
    "\n",
    "    ######################\n",
    "    # validate the model #\n",
    "    ######################\n",
    "    model.eval()\n",
    "    # calculate avg training accuracy over all users at every epoch\n",
    "    list_acc, list_loss = [], []\n",
    "    for client in range(num_users): # for each client\n",
    "        # get local model\n",
    "        local_model = LocalUpdate(dataset=train_dataset, idxs=user_groups[client],\n",
    "                                  gpu=gpu, optimizer=optimizer,\n",
    "                                  local_batch_size=local_batch_size, lr=lr,\n",
    "                                  local_epochs=local_epochs, loss_function=loss_function)\n",
    "\n",
    "        # get accuracy and loss of local model\n",
    "        acc, loss = local_model.inference(model=model)\n",
    "        list_acc.append(acc)\n",
    "        list_loss.append(loss)\n",
    "\n",
    "    # compute average accuracy\n",
    "    train_accuracy.append(sum(list_acc)/len(list_acc))\n",
    "\n",
    "    # print stats\n",
    "    print(f'\\nAverage training statistics (global epoch) : {epoch}')\n",
    "    print(f'|---- Trainig Loss : {np.mean(np.array(train_loss))}')\n",
    "    print('|---- Training Accuracy: {:.2f}% \\n'.format(100*train_accuracy[-1]))\n",
    "\n",
    "    if epoch in range(n_epochs - 9, n_epochs+1):\n",
    "        test_acc, test_loss = test_inference(model=model, test_dataset=test_dataset, gpu=gpu,\n",
    "                                             loss_function=loss_function)\n",
    "        test_acc_list.append(test_acc)\n",
    "        test_loss_list.append(test_loss)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# save the trained model\n",
    "\n",
    "filename_pt = 'fedAVG_results/{}_{}_{}_lr_[{}]_C[{}]_iid[{}]_unbalanced[{}]_E[{}]_B[{}]_{}_numGroups[{}]_percentage[{}].pt'\\\n",
    "    .format(\"ResNet50\", n_epochs, optimizer, lr, frac, iid, unbalanced,\n",
    "            local_epochs, local_batch_size, normalization_type, num_groups, percentage)\n",
    "torch.save(model.state_dict(), filename_pt)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# save the trained model\n",
    "\n",
    "filename_csv = 'fedAVG_results/{}_{}_{}_lr_[{}]_C[{}]_iid[{}]_unbalanced[{}]_E[{}]_B[{}]_{}_numGroups[{}]_percentage[{}].csv'\\\n",
    "    .format(\"ResNet50\", n_epochs, optimizer, lr, frac, iid, unbalanced,\n",
    "            local_epochs, local_batch_size, normalization_type, num_groups, percentage)\n",
    "\n",
    "data = list(zip(train_loss, train_accuracy))\n",
    "pd.DataFrame(data, columns=['train_loss','train_accuracy']).to_csv(filename_csv)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# test the trained model\n",
    "\n",
    "test_acc, test_loss = test_inference(model=model, test_dataset=test_dataset, gpu=gpu,\n",
    "                                     loss_function=loss_function)\n",
    "\n",
    "print(f'\\nResults after {n_epochs} global rounds of training:')\n",
    "print(\"|---- Avgerage Train Accuracy: {:.2f}%\".format(100*train_accuracy[-1]))\n",
    "print(\"|---- Test Accuracy: {:.2f}%\".format(100*test_acc))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "print(\"\\n\\n\")\n",
    "\n",
    "print(f'\\nResults after {n_epochs} global rounds of training:')\n",
    "print(\"|---- Test Loss: {:.2f}\".format((sum(test_loss_list) / len(test_loss_list))))\n",
    "print(\"|---- Test Accuracy: {:.2f}%\".format(100 * (sum(test_acc_list) / len(test_acc_list))))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}