fed_avg.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "GcvMbe-JowiY"
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import copy\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from util import view_10\n",
    "from data import fetch_dataset, data_to_tensor, iid_partition_loader, noniid_partition_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| using device: cuda\n"
     ]
    }
   ],
   "source": [
    "# set random seeds\n",
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "random.seed(0)\n",
    "\n",
    "# set device\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(\"| using device:\", device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "ZoZbOl8Gt3GV"
   },
   "outputs": [],
   "source": [
    "# hyperparams\n",
    "bsz = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VixfemygHq3q"
   },
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "IfPfvy-7sf-D"
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjwAAAD1CAYAAABUdy/PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAh5UlEQVR4nO3de7zNVf748fdCRXSIjFSDkksySCnJD1OkEUmmMC7RhWHo8h1GF1MKoaYaKWTcEvOTKZdompiQ6PJgZvR7qBTNRCdEcr+ck6zfH8xqrfW1t3322Xt/9lnn9Xw8zuPxXt77fPa7ts85y2fdlNZaAAAAQlYi6gIAAADSjQ4PAAAIHh0eAAAQPDo8AAAgeHR4AABA8OjwAACA4NHhAQAAwQuyw6OUWqGUOqKUOnDi67Ooa0LylFIVlVLzlVIHlVKblVK/iromFI5SqtaJe3RW1LUgOUqpgUqptUqpPKXUjKjrQeEopS5RSi1TSu1VSm1SSnWKuqZUC7LDc8JArXW5E191oi4GhfKCiOSLSBUR6S4iE5VSl0ZbEgrpBRFZE3URKJStIjJSRKZFXQgKRylVSkQWishiEakoIn1FZJZSqnakhaVYyB0eBEApVVZEOovI77XWB7TWq0TkdRHpGW1lSJZSqquI7BGRtyMuBYWgtZ6ntV4gIruirgWFVldEzhORZ7XWP2itl4nIagns52zIHZ7RSqlvlVKrlVKtoi4GSastIke11p9bf/aRiPCEpwhSSuWIyOMi8j9R1wIgLiUi9aMuIpVC7fAMFZGLROR8EZksIouUUjWjLQlJKici+7w/2ysiZ0VQCwpvhIhM1VrnRl0IAOMzEdkhIkOUUqcppa4XkZYicma0ZaVWkB0erfWHWuv9Wus8rfVLcvzRXLuo60JSDohIjvdnOSKyP4JaUAhKqUYi0lpEno24FAAWrfX3InKziNwoIttF5LciMldEgvqHSamoC8gQLccfz6Ho+VxESimlammtN574s4Yi8nGENSE5rUSkhohsUUqJHH96V1IpVU9r3TjCuoBiT2v9/+T4Ux0REVFKvSciL0VXUeoF94RHKVVBKdVWKVVaKVVKKdVdRFqIyN+irg0Fp7U+KCLzRORxpVRZpdQ1ItJRRF6OtjIkYbKI1BSRRie+JonIGyLSNrqSkKwTP19Li0hJOd5xLX1itQ+KIKVUgxOf4ZlKqcEiUlVEZkRcVkoF1+ERkdPk+FLJnSLyrYgMEpGbvUmvKFoGiEgZOT7G/H9FpL/Wmic8RYzW+pDWevt/v+T4cOURrfXOqGtDUoaJyGEReUBEepyIh0VaEQqjp4hsk+M/Z68TkTZa67xoS0otpbWOugYAAIC0CvEJDwAAgIMODwAACB4dHgAAEDw6PAAAIHh0eAAAQPDi7pmglGIJV8S01inbMJHPM3qp+jz5LKPHvRkW7s1wxPosecIDAACCR4cHAAAEjw4PAAAIHh0eAAAQPDo8AAAgeHR4AABA8OjwAACA4NHhAQAAwaPDAwAAgkeHBwAABI8ODwAACB4dHgAAELy4h4cC2ezyyy838cCBA51cr169TDxz5kwnN378eBP/85//TFN1AIBswhMeAAAQPDo8AAAgeEprHTupVOxkFilZsqSJy5cvn9D3+EMgZ555ponr1Knj5H7zm9+Y+A9/+IOT69atm4mPHDni5MaMGWPixx57LKG6fFprldQ3nkRR+TxjadSokdNetmyZiXNychK+zt69e01cqVKlQtdVEKn6PIv6Z5kO1113ndOePXu2iVu2bOnkPvvss0K/H/dm4Q0bNsxp2z8nS5Rw/z3eqlUrE7/zzjspr4V7MxyxPkue8AAAgODR4QEAAMGjwwMAAIKXVcvSq1WrZuLTTz/dyTVr1szEzZs3d3IVKlQwcefOnQtdR25urtN+7rnnTNypUycnt3//fhN/9NFHTi4d48zFzZVXXmni1157zcnZ87X8uWj255Kfn+/k7Hk7TZs2dXL2MnX/+0LQokULE/vzl+bPn5/pclKqSZMmTnvNmjURVYJ4evfubeKhQ4c6uWPHjsX8vnjzTYFE8IQHAAAEjw4PAAAIXqRDWvGWGSe6vDxV7Eep/lLJAwcOmNhe6ioism3bNhPv3r3byaVi6WtxYG8JICLSuHFjE8+aNcvEVatWTfiaGzduNPGTTz7p5ObMmWPi1atXOzn7sx89enTC71dU2Et7a9Wq5eSK4pCWvXT5wgsvdHLVq1c3sVIpW0GOQrI/l9KlS0dYSfF21VVXmbhHjx4m9rdwuPTSS2NeY/DgwSbeunWrk7Onntg/x0VEPvzww4IVmyI84QEAAMGjwwMAAIJHhwcAAAQv0jk8W7Zscdq7du0ycSrm8PjjhHv27DHxz3/+cydnL0F++eWXC/3eSNyLL77otO3jOpJlzwMqV66ck7O3C7DntIiINGjQoNDvnc3sU+Tff//9CCtJDXte19133+3k7HkDGzZsyFhNcLVu3dppDxo0KOZr7c+pffv2Tu6bb75JbWHFTJcuXZz2uHHjTHzOOeeY2J/vtmLFChNXrlzZyT311FMx38++jv99Xbt2PXXBacATHgAAEDw6PAAAIHiRDml99913TnvIkCEm9h9n/utf/zKxvfOxb926dSZu06aNkzt48KCJ/aV2995776kLRspcfvnlJr7xxhudXKwlxP7O1YsWLTKxf4q9vUTS/rsj4m4fcO211yb03qHwT6Au6qZMmRIzZ29NgMyylyRPnz7dycWbrmAPkWzevDn1hQWuVCn3V/oVV1xh4j/96U9Ozt4OZOXKlSYeMWKE87pVq1aZ+IwzznByc+fONfH1118fs661a9fGKztjwvrpBwAAcBJ0eAAAQPDo8AAAgOBl1WnpCxYsMLF9zISIe/p1w4YNndydd95pYnsuhz1nx/fxxx877b59+xaoVhSMf4zI0qVLTZyTk+Pk7FOR33zzTRP7y9XtLdD940DsuR07d+50cvap9v7pzPZ8Intpu4h7knpR4S+zr1KlSkSVpEe8+SD23zFk1u23327i8847L+br7CXPIiIzZ85MV0nFgn1EhEj8OW72/WEvWd+3b1/M7/GXtsebt5Obm2vil156KebrMoknPAAAIHh0eAAAQPCyakjLFu+x2t69e2Pm7N1WX3nlFSfnD18gvWrXrm1ie8sBEXco4ttvv3Vy9gn09qNQ+9R6EZE33njjpHFhlClTxsS//e1vnVz37t1T8h6Z1K5dO6dt//cVRf6QnH9Cuu3rr79Odzk4wd6pV0TkjjvuMLH/c9fe8X7kyJFpras4sJeRP/TQQ07Onh4wYcIEJ2dPA4j3+9b28MMPJ1zXPffcY2J/WkFUeMIDAACCR4cHAAAEjw4PAAAIXtbO4Yln+PDhTts+psBequyf0rtkyZK01lXc+duO21sE+HNJ7G0G7BO8RdxtyKOcc1KtWrXI3jtV6tSpEzPnb81QFPhHiNhzej7//HMnZ/8dQ+rVqFHDxK+99lrC3zd+/HgTL1++PJUlFQuPPPKI07bn7eTn5zu5t956y8RDhw51cocPHz7p9UuXLu207aXn/s9E+ygefz7WwoULT3r9KPGEBwAABI8ODwAACF6RHNLyd1C2l6Lbu+H6p8Paj0/901tfeOEFE9tL+ZC4yy67zGn7w1i2jh07mtg/BR2ZsWbNmqhLMOzdtm+44QYnZ+8eG29nV/+UZ3v5M1LP/pz8Hb1tb7/9ttMeN25c2moKVYUKFUw8YMAAJ2f/vrKHsEREbr755oSuf/HFF5t49uzZTs6eMuJ79dVXTfzkk08m9F5R4gkPAAAIHh0eAAAQvCI5pOX74osvTNy7d28TT58+3Xldz549TxqLiJQtW9bE/gF29s6/iO2ZZ55x2vYMfn/YKluGsUqUcPv8xWk37ooVKyb1ff7hvfbn7K+MvOCCC0x8+umnm9jftdr+HPzVIx9++KGJ8/LynFypUj/+CPvHP/5xytpROPYQyZgxY2K+btWqVSa2DxIVib9TPk7Ovnf8Xa1t9u7GIiI/+clPTNynTx8nd9NNN5m4fv36Ji5XrpzzOnvIzJ/uMWvWLBPHO6w7W/CEBwAABI8ODwAACB4dHgAAELwg5vDY5s+fb+KNGzc6OXuOyXXXXefknnjiCRNXr17dyY0aNcrEnMDsat++vYkbNWrk5Ozx3tdffz1TJRWIP2fHrnndunUZrib1/Pkw9n/fpEmTnJx/0nIs/hJkew7P0aNHndyhQ4dM/Mknn5h42rRpzuvsbSL8+V3ffPONiXNzc52cvRP3hg0bTlk7CsbeTVkk8R2V//3vf5vY/vyQHHsHZf/k8cqVK5v4P//5j5NLdIuVrVu3mtg/Ob1q1aom/vbbb53cokWLErp+tuAJDwAACB4dHgAAELzghrRs69evd9q33XabiTt06ODk7CXs/fr1c3K1atUycZs2bVJZYpFnDynYSydFRHbs2GHiV155JWM1+fxDTf3DZ23Lli0z8YMPPpiukjLG35V18+bNJm7WrFlS19yyZYvTXrBggYk//fRTJ/fBBx8k9R62vn37mth+fC/iDp0g9fwDJxPdtiHeknUUnL1ruL978uLFi03sbzVhb9niH+Y5Y8YME3/33XcmnjNnjvM6e0jLzxU1POEBAADBo8MDAACCR4cHAAAEL+g5PD57HPTll192clOmTDGxvV29iEiLFi1M3KpVKye3YsWKlNUXGvsYgEwfz2HP2xk2bJiTGzJkiIn9Zc5PP/20iQ8cOJCm6qIzduzYqEsoMH8LCVuiy6SROHt7iXin09v8+SGfffZZKkuCxT5qReR/z2tLhv07rmXLlk7OnrdV1OfM8YQHAAAEjw4PAAAIXtBDWv6OsL/85S9N3KRJEyfnD2PZ7B1iV65cmaLqwpfJ3ZX9XZ7tYasuXbo4Ofvxe+fOndNaF9LL3lkdqbFkyRITn3322TFfZ2850Lt373SWhDSztxeJt/s8y9IBAACyHB0eAAAQPDo8AAAgeEHM4alTp46JBw4caOJbbrnFed25556b0PV++OEHp20vqU50a/Xiwj4p245F3C3Q77333pS/9/3332/i3//+906ufPnyJp49e7aT69WrV8prAUJRqVIlE8f7eTdhwgQTh7iFQ3Hy1ltvRV1CRvCEBwAABI8ODwAACF6RGdKyh6O6devm5OxhrBo1aiR1/bVr15p41KhRTi6Ty6uLGnvJoh2LuJ/Zc8895+SmTZtm4l27djm5pk2bmrhnz54mbtiwofO6Cy64wMT+Cd72I1r70TuKNn/YtHbt2iZOxcnsxdH06dOddokSif07+L333ktHOYhA27Ztoy4hI3jCAwAAgkeHBwAABI8ODwAACF5WzeGpUqWKievVq+fknn/+eRPXrVs3qevbp8w+9dRTTs4+boCl56lRsmRJEw8YMMDJ2Uc67Nu3z8nVqlUroevbcwiWL1/u5B555JGE60TR4c8TS3S+CVz2USytW7d2cvbPv/z8fCf3wgsvmPibb75JT3HIuIsuuijqEjKCnxYAACB4dHgAAEDwMj6kVbFiRRO/+OKLTs5+zJrsIzZ7mOPpp592cvZS5cOHDyd1fbjef/99E69Zs8bJ+SfS2+wl6/ZQps9esu6f1JuO3ZtRtFx99dUmnjFjRnSFFDEVKlQwcbwd6L/++munPXjw4HSVhAi9++67JvaHiUOa4sETHgAAEDw6PAAAIHh0eAAAQPDSMofnqquuMvGQIUOc3JVXXmni888/P6nrHzp0yGnbxxY88cQTJj548GBS10ficnNzTeyfTt+vXz8TDxs2LOFrjhs3zsQTJ0408aZNm5IpEQHxj5YAUHjr16838caNG52cPZ+2Zs2aTm7nzp3pLSzFeMIDAACCR4cHAAAELy1DWp06dTppfCqffPKJiRcvXuzkjh49amJ/ufmePXsKWCHSYdu2bU57+PDhJ42BgnjzzTdNfOutt0ZYSTg2bNhgYv/U8+bNm2e6HGQRe1qIiMiUKVNMPGrUKCc3aNAgE9u/v7MVT3gAAEDw6PAAAIDg0eEBAADBU/7pw05SqdhJZITWOmXrcPk8o5eqz5PPMnrcm2Hh3jwuJyfHac+dO9fErVu3dnLz5s0zcZ8+fZxclNvCxPosecIDAACCR4cHAAAEjyGtLMdj87Dw2Dwc3Jth4d48OXuIy1+W3r9/fxM3aNDAyUW5TJ0hLQAAUGzR4QEAAMGjwwMAAILHHJ4sxzyBsDBPIBzcm2Hh3gwHc3gAAECxRYcHAAAEL+6QFgAAQAh4wgMAAIJHhwcAAASPDg8AAAgeHR4AABA8OjwAACB4dHgAAEDw6PAAAIDgBdfhUUqdoZSaqpTarJTar5Rap5T6RdR1ITlKqYFKqbVKqTyl1Iyo60HhKKVmKaW2KaX2KaU+V0rdFXVNSA73ZpiUUrWUUkeUUrOiriXVSkVdQBqUEpGvRKSliGwRkXYiMlcp9TOt9ZdRFoakbBWRkSLSVkTKRFwLCm+0iNyptc5TStUVkRVKqX9prf8RdWEoMO7NML0gImuiLiIdgnvCo7U+qLUerrX+Umt9TGu9WET+IyKXR10bCk5rPU9rvUBEdkVdCwpPa/2x1jrvv80TXzUjLAlJ4t4Mj1Kqq4jsEZG3Iy4lLYLr8PiUUlVEpLaIfBx1LQBElFITlFKHRGSDiGwTkb9GXBJQ7CmlckTkcRH5n6hrSZegOzxKqdNEZLaIvKS13hB1PQBEtNYDROQsEfk/IjJPRPLifweADBghIlO11rlRF5IuwXZ4lFIlRORlEckXkYERlwPAorX+QWu9SkQuEJH+UdcDFGdKqUYi0lpEno24lLQKcdKyKKWUiEwVkSoi0k5r/X3EJQE4uVLCHB4gaq1EpIaIbDn+61PKiUhJpVQ9rXXjCOtKqVCf8EwUkUtEpIPW+nDUxSB5SqlSSqnSIlJSjt+ApZVSQXbUQ6eU+olSqqtSqpxSqqRSqq2IdJNAJ0iGjnszKJPl+D88Gp34miQib8jxFXjBCK7Do5SqLiL95PiHtl0pdeDEV/doK0OShonIYRF5QER6nIiHRVoRkqXl+PBVrojsFpE/iMh9WuvXI60KyeLeDITW+pDWevt/v0TkgIgc0VrvjLq2VFJa66hrAAAASKvgnvAAAAD46PAAAIDg0eEBAADBo8MDAACCR4cHAAAEL+6eCUoplnBFTGutUnUtPs/operz5LOMHvdmWLg3wxHrs+QJDwAACB4dHgAAEDw6PAAAIHh0eAAAQPDo8AAAgODR4QEAAMGjwwMAAIJHhwcAAASPDg8AAAgeHR4AABA8OjwAACB4dHgAAEDw4h4eCkRh3LhxJr7nnntMvH79eud17du3N/HmzZvTXxgAIKXefvttp63Uj+d+XnvttSl9L57wAACA4NHhAQAAwStWQ1pnnXWWicuVK+fkbrzxRhNXrlzZyT3zzDMmzsvLS1N1xVeNGjWcdo8ePUx87NgxE19yySXO6+rWrWtihrSyQ+3atZ32aaedZuIWLVqYeMKECc7r7M85WQsXLnTaXbt2NXF+fn6hrw/382zWrJmJn3jiCed111xzTcZqQtHz7LPPmtj+eyQiMnPmzLS9L094AABA8OjwAACA4NHhAQAAwQtuDo89H2To0KFO7uqrrzZx/fr1E75m1apVTWwvk0Zq7Ny502mvXLnSxDfddFOmy8EpXHrppU67d+/eJr711ludXIkSP/6b6rzzzjOxP2dHa13ouvy/K5MmTTLxfffd5+T27dtX6PcrjsqXL2/i5cuXm3j79u3O684999yYORQ/Y8aMcdq//vWvTfz99987OX+ZeirxhAcAAASPDg8AAAhekRzSspcji7iPq7t3727iMmXKOK+zd3D86quvnNz+/ftN7C9/vu2220zsL6fdsGFDglUjloMHDzptlphnt9GjRzvtdu3aRVRJfL169TLx1KlTndzq1aszXU7Q7CEsv82QFpo2beq07e0NVq1a5eTmzp2btjp4wgMAAIJHhwcAAASPDg8AAAhe1s7hsZc/ioiMHTvWxF26dHFy9pER8WzcuNHEbdu2dXL2mKI/L+ecc845aYzUqFChgtNu2LBhNIUgIUuXLnXa8ebw7Nixw8T2PBp7ubpI/KMl7K3nW7ZsmXCdyBx7fiSKBvuol4cfftjE3bp1c1733XffJXV9+zr+NjBffPGFiQcPHpzU9ZPBEx4AABA8OjwAACB4WTuk1alTJ6d91113Ffga9mMzEZE2bdqY2F+WfvHFFxf4+kiNM88802lXq1Ytoe9r0qSJif1hSJa2p8/EiROd9oIFC2K+1t5FNdnlyTk5OSZev369k7N3b/bZda1duzap90Zi/J2yS5cuHVElSNTkyZNNXKtWLRPXq1fPeZ2/bDxRDz30kIkrVark5O6++24Tf/TRR0ldPxk84QEAAMGjwwMAAIJHhwcAAAQva+fw+Kcux/Pll1+aeM2aNSb2T0v35+3Y/OMkkDlbt2512jNmzDDx8OHDY36fnduzZ4+Te/7551NQGU7m6NGjTjvefZUK9hYSZ599dsLfl5uba+K8vLyU1oT4rrjiChN/8MEHEVaCWA4dOmRiew5WsvOvGjVq5LSrV69uYn/biajmePGEBwAABI8ODwAACF7WDmnZy9ZERPr27WviJUuWOLlNmzaZ2N7ZtSCqVKmS1Pch9UaMGGHieENaCFPXrl2dtv2zoEyZMglf55FHHklZTTjOHs7cu3evif2d8WvWrJmxmpAY++eqiMjPfvYzE3/66acmLsgy8bJly5rYn0JibzfiD2u++uqrCb9HKvGEBwAABI8ODwAACB4dHgAAELysncPjL1VO91yOq6++Oq3XR3LsU7XjnaiNoqV79+5O+4EHHjCxf8zLaaedltA1161b57TtYy2QGvb2D++++66J27dvH0E1OJWf/vSnJvbnxdrzsQYOHGjinTt3Jnz9Z555xsT+VjL27/Brrrkm4WumE094AABA8OjwAACA4GXtkFay7rnnHhPbS+ZOxV6i53vvvfdM/P777ydXGJJiD2P5JzIjGjVq1HDaPXv2NHHr1q0Tukbz5s2ddqKf7b59+5y2PRT217/+1ckdPnw4oWsCoahfv77Tnj9/vonPOeccJzd+/HgTv/POOwldf/DgwU67d+/eMV87atSohK6ZSTzhAQAAwaPDAwAAgldkhrTsXRvr1avn5B599FETt2vXLuY1El3x468Q69Onj4l/+OGHUxcLBMZ+VP766687uWrVqmWsDntlkIjI5MmTM/beSFylSpWiLiFYpUq5v7Z79Ohh4qlTpzq5eL/z7JXJDz74oIntlVciIhUrVjSxvxJLKWXimTNnOrkXX3zx5P8BEeIJDwAACB4dHgAAEDw6PAAAIHhZNYfH3lH1sssuc3KvvfaaiatWrerk7OWn9vwbfwn5DTfcYGJ7TpDPHyO95ZZbTDxu3Dgnl5+fH/M6QIjscfuTtRNhzy0QSXwXbX9H31/84hcmfvPNNwtcB9LjpptuirqEYHXt2tVpT5kyxcT+9g72fbVp0yYnd8UVV5w07tixo/O6888/38T+7157V+Y77rjjlLVHjSc8AAAgeHR4AABA8CId0jr99NOdtj3kNG/evJjf99hjjzntZcuWmXj16tUmtpfT+a/zd6S0Va5c2WmPHj3axFu2bHFyCxYsMHFeXl7MayI5iW4l0KJFC6f9/PPPp62m4mj9+vUmbtWqlZOzl8W+9dZbTu7IkSMFfq8777zTaQ8aNKjA10D6LV++3MQcHppeXbp0MfH06dOdnH1Irn24q4jIr371KxPv3r3byT399NMmbtmypYnt4S0Rd8jaHzKzd2/+6quvnJz9c+KLL76QbMATHgAAEDw6PAAAIHh0eAAAQPBUvFOKlVIpP57aXnr++OOPO7khQ4bE/D57yal9OrOIO25pz7/xT09u3Lixif3l5E8++aSJ/fk9/jI929///ncTjx071sn5Y6a2devWxczZtNYFX/MbQzo+z3Szj/IoyGnpDRo0MPEnn3yS0poKI1WfZ1H8LBNVvnx5p71r166Yr+3QoYOJM70svbjfm507dzbxX/7yFydnbxXiHwW0efPm9BaWpGy+N+35p9WrV3dyI0eONLE/vyce+3Oxj4Gwj5wQiT+Hx/bnP//Zaffq1SvhWlIt1mfJEx4AABA8OjwAACB4aV+WXrJkSac9YsQIEw8ePNjJHTx40MQPPPCAk5szZ46J/aV39jI6ezmyv1vzxo0bTdy/f38nZy+xzMnJcXLNmjUzcffu3Z2cvaPo0qVLJRZ/yd6FF14Y87X40aRJk0zcr1+/hL+vb9++Jr7vvvtSWRLSrG3btlGXgAQcPXo0Zs4eBjnjjDMyUU7QFi5caGJ/yxb/d0ui7CXl8bZp6datm4nt7Sl8ubm5SdWRSTzhAQAAwaPDAwAAgkeHBwAABC/tc3jsuRQi7rydQ4cOOTl7jsaSJUucXNOmTU3cp08fJ2efmFymTBkT+8ve7SV78cY99+3b57T/9re/nTQWccc37W28fffff3/MHGLbsGFD1CUUG/aWEddff72Ts5fF2kuOU8W+p8eNG5fy6yP17Hkl/n1at25dE/tz6AYMGJDWukKUinvC3+7h1ltvNbE9b9U/BmLu3LmFfu9swRMeAAAQPDo8AAAgeGnfaXnbtm1O294J2T9d3H4sWrZsWSd38cUXJ/R+w4cPN7F9yrmIu2tvUVHcd3O1ff755067Zs2aMV9rn7Lu/92J8uTebNrNtXnz5k774YcfNnGbNm2cnL2NQrLLYCtWrGjidu3aObnx48eb+Kyzzop5DX84zd4Wwt5aIhO4N3/0xz/+0WnbQ5RVqlRxckeOHMlESQWWTfdmOjz44INO294iZufOnSZu0qSJ87qisNzcx07LAACg2KLDAwAAgkeHBwAABC/ty9K3b9/utO05PP6W4w0bNox5Hfvk85UrVzq5BQsWmPjLL780cVGcs4PYPv74Y6d90UUXxXztsWPH0l1OkWcfwyISf3v53/3udybev39/Uu9nzwtq3Lixk4s3l3DFihUmnjhxopPL9LwdJMb+PPPz8yOspHizT1a/6667nJz9GU2ePNnERXHOTqJ4wgMAAIJHhwcAAAQv7UNaLVq0cNo333yzif3H2jt27DDxtGnTnNzu3btNzCPS4sl+7Coi0qFDh4gqKX769++f1uvb9/6iRYuc3L333mvibF3SDJe9c2/Hjh2d3Pz58zNdTrG1dOlSE9vDWyIis2bNMvGjjz6asZqixBMeAAAQPDo8AAAgeHR4AABA8NJ+tAQKh+3rf+SPQS9evNjEl1xyiZNT6sf/bbVr13ZyHC1xXKNGjZz2oEGDTHz77bcX9vIi4v6/PnTokInfffdd53X2/Kz169en5L3TjXvzR1u3bnXaZ599tokvu+wyJ+efrJ4tsuneTBX7OAn7KAkR97T00OZVcbQEAAAotujwAACA4DGkleV4bB6WbH5sbu983rt3byc3cuRIE9vDFSLuTuf2MlgRkYULF5rY33W9qOPe/NGcOXOctj3EbJ9oLyKyefPmjNRUUNl8b6JgGNICAADFFh0eAAAQPDo8AAAgeMzhyXLMEwgL8wTCwb0ZFu7NcDCHBwAAFFt0eAAAQPDo8AAAgODR4QEAAMGjwwMAAIJHhwcAAASPDg8AAAgeHR4AABA8OjwAACB4cXdaBgAACAFPeAAAQPDo8AAAgODR4QEAAMGjwwMAAIJHhwcAAASPDg8AAAje/wdBSLxIvevFWgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 720x288 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_data, test_data = fetch_dataset()\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size = 1000, shuffle=False) # inference bsz=1000\n",
    "debug_loader = torch.utils.data.DataLoader(train_data, bsz)\n",
    "img, label = next(iter(debug_loader))\n",
    "view_10(img, label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "a4mN1wUr4izX"
   },
   "outputs": [],
   "source": [
    "# get client dataloaders\n",
    "iid_client_train_loader = iid_partition_loader(train_data, bsz = bsz)\n",
    "noniid_client_train_loader = noniid_partition_loader(train_data, bsz = bsz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "ls_0Ke21Govu"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iid:  tensor([59., 67., 65., 61., 55., 52., 55., 64., 48., 74.])\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjwAAAD1CAYAAABUdy/PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAgJUlEQVR4nO3de9TNVf7A8c92T8Lzi9xlxVCTCc3KSCylC8WkGZHpRoYGZcYtCkU8sqxxacbqOpmmQa6pKUolkxl3EpMuy0hyS0XKnQf798djdnt/Ocd5zvM95/ucfd6vtaz12c/+nnM+9fU9z7avSmstAAAAPisWdQIAAACpRoMHAAB4jwYPAADwHg0eAADgPRo8AADAezR4AACA92jwAAAA73nZ4FFKTVNKfaWU2q+U2qSU6hF1TkieUur/lFKvKqUOKaW+VErdGXVOSI5S6kGl1Fql1DGl1N+izgfJU0odDPw5qZSaHHVeSE42PJslok4gRcaKyG+11seUUpeKyPtKqQ+11h9EnRiS8pSIHBeRKiLSWEQWKKU2aK0/jjQrJGOXiOSKSBsROS/iXFAIWuty/4uVUuVEZLeIzIkuIxSS98+mlz08WuuPtdbH/lc8/aduhCkhSUqp80Wko4g8qrU+qLVeKiKvi8g90WaGZGit52mtXxORvVHnglB1FJFvROTfUSeC5GTDs+llg0dERCn1tFLqsIh8JiJficibEaeE5NQXkRNa603WzzaIyOUR5QPgTF1F5O+as4pQhHnb4NFa9xGRC0SkpYjME5Fj8V+BIqqciOwP/OwHyb+3ACKmlLpYRFqJyEtR5wLE422DR0REa33y9BBITRHpHXU+SMpBESkf+Fl5ETkQQS4AznSPiCzVWn8RdSJAPF43eCwlhDk8mWqTiJRQSv3E+lkjEWHCMlA03Cv07iADeNfgUUpdpJTqopQqp5QqrpRqIyK/EZH3os4NBae1PiT5Q5KjlFLnK6WuEZEOIjI12syQDKVUCaVUGREpLiLFlVJllFK+rhb1nlKquYjUEFZnZbxseDa9a/BI/oqs3iKyQ0T2ich4EemntX490qxQGH0kf5nkNyIyQ0R6syQ9Yw0XkSMi8rCI3H06Hh5pRiiMriIyT2vNEHPm8/7ZVEyqBwAAvvOxhwcAAMBBgwcAAHiPBg8AAPAeDR4AAOA9GjwAAMB7cdfYK6VYwhUxrbUK6724n9EL635yL6PHs+kXnk1/xLqX9PAAAADv0eABAADeo8EDAAC8R4MHAAB4jwYPAADwHg0eAADgPRo8AADAezR4AACA9+JuPFiU3HXXXSYeOXKkU1evXj0Ta/3jnk+fffaZc92vf/3rmHUAAMBf9PAAAADv0eABAADeo8EDAAC8lzFzeMqXL2/iSy65xKk7derUWV/ToEEDp7xmzRoT5+bmOnWTJk0y8fHjx5POE4kpU6aMiRcsWODUtW7d2sT2nKxXXnnFue7xxx838aeffurUnTx5MpQ8kRoTJkxwygMGDDCxfV9FRJ577jkTf/XVV6lNDEBMd999t1OeOnVqzGvt39kHDhxIWU4FQQ8PAADwHg0eAADgPWUPGZxRqVTsyjRTSpn4vPPOS+g1HTt2dMr2sFVOTo5T1759exO/9dZbyaSYElprde6rEhPl/bzggguc8qxZs0zcpk0bp84ejtq1a5eJa9WqFfP9g8MgTzzxhIlPnDhRsGRTKKz7WZSezUSVLVvWxDNnznTq2rVrZ+Lgd5L9PP7yl79MUXYF58uziXzZ/Gwm6uWXX3bKXbp0MfHEiROduoceesjE8doZqRDrXtLDAwAAvEeDBwAAeI8GDwAA8F7GLEu3xwAPHz7s1NWtW9fEn3/+uYmDS+Z2795t4oULFzp1Q4cONXFRmsPjC3uOlMiZ83Zsf/nLX0z8yCOPmLhbt27OdaNHjzbxiBEjnLqvv/7axPayZkSnXLlyJm7SpEnCr6tYsWIKsgGQiAoVKpi4evXqTp29JcyHH37o1KV73k4i6OEBAADeo8EDAAC8lzHL0uOxl6kfOXIk5nWVK1c28apVq5w6u9u8adOmTt3mzZsLmWHyfFn6OmzYMKc8atQoEweXjderV8/E27dvj/me//73v03cvHlzp27Hjh0mtoc8z/Z56cTS13zjx493yv379zdx8DtpxYoVJm7ZsmVqEysAX57N4HYP27ZtM3FwqfHAgQPTklMUeDbPrnHjxiZet26dU3fs2DETJ7pdTDqwLB0AAGQtGjwAAMB7NHgAAID3MmZZejzx5u3Y6tevb+LgUQf2PJ39+/eHk1iWq1atmol79OgR87opU6Y45Xjzdmz2CenBOTw1a9Y08fLly526O++808RRzs/KNvY8uVatWkWXCBzxnjf7FHsRkU6dOpl40KBBJp49e3b4iSES9hEwIvGP9Mk09PAAAADv0eABAADe82JZeiylS5d2yvbJ28HdW+2dlseNG5fSvAoik5e+2rsk5+bmxryuRo0aTtneETseexnktdde69TZXezBLtrvv//exJMnT3bqRo4cmdBnJyubl77aXeNbt2516ooV+/HfXvburSLukCTL0sPXuXNnpzxr1qxCv6e9lcCcOXOcurlz55o40eHrdMjmZ9N24403OuW333475rX29I+itCM6y9IBAEDWosEDAAC8R4MHAAB4z4tl6fYckKuvvtrE9rJJEXeMccOGDU5d8PR0FN4NN9yQ0HUHDhxI6v3t7QiCJ9xXqVLFxPbJ6SLu34OePXs6dZ9//rmJp06dmlReOLfg3EF73k6wbvTo0WnJCQVjHzthf+8Gy8E6e2n7HXfc4dQVpTk92apt27Yx64JbtiT6HV9U0MMDAAC8R4MHAAB4r8gOadWpU8cpz5w508TBXZJzcnJMbA9lBO3bt8/EY8aMceo2btyYTJpI0tKlS018/Pjx0N//8OHDJr7wwgudukWLFpn4mmuucers5fNLlixx6uxTpFFwF198cVKv69q1q4nfeeedsNLBabfffntSr0v05PQJEyY4ZXv35mBdcIk80qNBgwYmjncPHnzwQae8du3alOWUCvTwAAAA79HgAQAA3qPBAwAAvFdk5/BUr17dKTdq1MjEpUqVSuo97bk+wSXr//3vf038n//8J6n3R+Ls/8d5eXkp/azgHKHFixebODiHxz5l/aKLLnLqmMNTOH379k3qdbVr1w45EyQreExEIoJzfXbs2GFie2m7iEj//v1NPGnSpAJ/FpLTo0cPEweP+rHnvm7ZsiVtOaUCPTwAAMB7NHgAAID3iuyQln1CsojIVVddZeLgKeiJat++vYmD3etPPvmkiceOHevUvfvuu0l9XrYJbhdQoUKFiDKJb9y4cSYOnrJelE7j9oG9vcQVV1yR1Hu88cYbIWWDsynI7sZh7IRsD1XZuy6LuENcDGmlT5cuXWLW2b//gr+XMw09PAAAwHs0eAAAgPeK7JBWUBg7Ia9bt87EH3/8sVNnd58+88wzTt0f/vAHEy9YsKDQefjqsssuc8pNmjSJKJP47ENHd+3aFWEm/qtataqJ69evH/O6YsV+/LdXnz59nLrg84hwrVq1KiXXJsJesRVUq1Ytp8zBotF44YUXok4hNPTwAAAA79HgAQAA3qPBAwAAvJcxc3jCoLU28dy5c506eyx52bJlTp19gjZzeMIxZcqUyD67WrVqJr755psjyyPb2M9f0KlTp0z8+uuvpyMdnLZixYqEr03nbuPM2UmtcuXKmbhEiR+bAh999JFzXRjzZ4sKengAAID3aPAAAADvZdWQVjyrV6828WuvvebUdejQwcS33nqrU0f3e3L27t0b2Wf369fPxOXLl48sD6AoCA4d2QeEBoe7Vq5cWejPs5ebB3daLsjwGs5NKWXi4O+uAQMGmLhKlSomDh6evXv37hRll3708AAAAO/R4AEAAN6jwQMAALzHHJ7T7GWxf//735262267zcQlS5ZMV0oZZ82aNU55yZIlJm7VqpVT16JFCxPPmDEjpXnZ49giIhUrVox57fvvv2/iDz74IEUZZY977723wK8J3p+dO3eGlA0S0blz55S+/6xZs2LWPfnkkyn97GxTunRpEw8cONCps7+DswU9PAAAwHs0eAAAgPciHdIaMWKEU77rrrtMHFyuuGHDhpTmUrx4cRN3797dqTt8+LCJP/3005TmkcmCO+nG21nXPhE7uOt1Xl5eqHnZOyuLiPTo0SPmZ73xxhsmjpc/EtOoUaMCv8ZeFi3i7ob95ZdfFjonpFdwiOzqq682cfBez549Oy05ZQt7qsaBAwdiXtetWzcTB7+PfUIPDwAA8B4NHgAA4D0aPAAAwHuRzuGpU6eOU65bt66J33rrLafu/vvvN/H8+fNDz+Wxxx4zcfv27Z26PXv2mPiTTz4J/bOzkX06/cmTJ1P6WW+++WbMuuD9ZFlsuOwtAYLbA9iKFSuW0HXIDPbxEePHj4953cSJE9ORTtayn6VSpUrFvG7IkCEmXrRokVNnz2HNdPTwAAAA79HgAQAA3ot0SGvSpElO+Ve/+pWJ7dNbRUReffVVE48aNcqps7tFjx49auIyZco41+Xk5Jh4+vTpTl3z5s1NHBxisbv7kLht27bFrLOXqtrLxEVEDh06VODPCi49v+eee0zcsGFDp27Tpk0mnjx5coE/C4mzl/bHW+b/2Wefmdhehi7CUvRMNGHCBBPbw1si7indYZy+jtguu+wyE19//fUxr1u/fr2Jd+3alcqUIkUPDwAA8B4NHgAA4D0aPAAAwHsq3ri6Uiqte+tfc801Jn7++eeduksvvTTm6+z/Bns5+y233JLwZ9vzTYJHXgRPT08nrXVoa3TTfT8rVKhg4qlTpzp17dq1M/Gf//xnp65///4JvX/16tVN/PDDDzt1DzzwgIlPnDjh1HXo0MHECxcuTOizwhLW/Uz3vUzUhRde6JQXL15s4ssvvzzm6+wtKTJlzk4mP5thCx4fYZ+Ivn37dqeudu3aacmpoHx8Np999lkT21u7BM2YMcPE9hFPmSrWvaSHBwAAeI8GDwAA8F6ky9KDli1bZmL7RF0RkWHDhpnYXr4u4naHxxvGst9/zJgxTt2HH35o4m+++SbBjBHPDz/8YOLc3Fyn7tprrzXx73//e6fuuuuuM7HdNR48xb5s2bImrlq1qlO3ceNGE3fs2NGp27x587lSR5Jat27tlGMNY7300ktOOVOGsXB28XZTDg53IX0S3eLD/p71GT08AADAezR4AACA92jwAAAA7xWpZek4k69LX+15Om3btnXqBg0alNB72HNx5s2b59TZW9vbp91Hzcelr7ZOnTo5ZXu5qy14jzPxlHpfn81ENWvWzMQrVqxw6uzjfgYOHJi2nArDx2ezdOnSJp4/f75TZx81YR//sXPnztQnlmIsSwcAAFmLBg8AAPAeQ1pFXLZ3m/vGx27zbJXtz+bs2bNNHBzKtHdTDu60XFTxbPqDIS0AAJC1aPAAAADv0eABAADeYw5PEZft8wR8wzwBf2Tjs2kvX962bZuJg8vSmzdvnracwsKz6Q/m8AAAgKxFgwcAAHivSJ2WDgDIDPZy8wEDBkSYCZAYengAAID3aPAAAADv0eABAADeY1l6EZeNS199xtJXf/Bs+oVn0x8sSwcAAFmLBg8AAPBe3CEtAAAAH9DDAwAAvEeDBwAAeI8GDwAA8B4NHgAA4D0aPAAAwHs0eAAAgPdo8AAAAO952eBRSv2fUupVpdQhpdSXSqk7o84JyVFKPaiUWquUOqaU+lvU+SAcSqmfKKWOKqWmRZ0LkqOUKq2UmnL6O/aAUmq9UurmqPNCcrLhu7ZE1AmkyFMiclxEqohIYxFZoJTaoLX+ONKskIxdIpIrIm1E5LyIc0F4nhKRNVEngUIpISLbRaSViGwTkVtEZLZS6mda661RJoakeP9d610Pj1LqfBHpKCKPaq0Paq2XisjrInJPtJkhGVrreVrr10Rkb9S5IBxKqS4i8r2IvBdxKigErfUhrfVIrfVWrfUprfV8EflCRH4edW4ouGz4rvWuwSMi9UXkhNZ6k/WzDSJyeUT5ADhNKVVeREaJyICoc0G4lFJVJP/7l550FEk+NnjKicj+wM9+EJELIsgFgGu0iEzRWu+IOhGERylVUkSmi8hLWuvPos4HOBsf5/AcFJHygZ+VF5EDEeQC4DSlVGMRuUFEmkScCkKklComIlMlf97kgxGnA8TkY4Nnk4iUUEr9RGv939M/ayR0swJRu1ZE6ojINqWUSH5vbHGl1E+11ldGmBeSpPJv5BTJXyByi9Y6L+KUgJi8a/BorQ8ppeaJyCilVA/JX6XVQUSaR5oYkqKUKiH5f0+LS/4vxzKSP0frRLSZIQnPi8hMqzxI8htAvSPJBmF4RkQuE5EbtNZHok4GycuG71of5/CIiPSR/GV134jIDBHpzZL0jDVcRI6IyMMicvfpeHikGSEpWuvDWuvd//sj+cPPR7XW30adGwpOKXWxiPxO8v9RuVspdfD0n7uizQxJ8v67Vmmto84BAAAgpXzt4QEAADBo8AAAAO/R4AEAAN6jwQMAALxHgwcAAHgv7j48SimWcEVMa63Cei/uZ/TCup/cy+jxbPqFZ9Mfse4lPTwAAMB7NHgAAID3aPAAAADv0eABAADeo8EDAAC8R4MHAAB4jwYPAADwHg0eAADgPRo8AADAezR4AACA92jwAAAA79HgAQAA3ot7eChQlCxfvtwpz5kzx8STJk1KdzpAVqtTp46Ju3fv7tR17drVxLVr1074PadNm2bie++916nTmjM5UTj08AAAAO/R4AEAAN5T8boJlVIZ14dYrNiPbbjc3Fyn7pFHHon5umeffdbE3377bVKfPXbsWKd85MiRpN7HprVWhX6T0zLxftqCf1cnTpxo4oEDB6Y7naSEdT+jvJclSrgj4TVr1jTx1q1b05xNYsqUKeOUb7rpJhPPmDHDqfvuu+9M/OKLLzp1jz32mImz4dksXry4U7aHmYYOHWriSy65JPTPbtCggVPevHlz6J9h8+HZTIWePXuauFKlSk5d8HdeURHrXtLDAwAAvEeDBwAAeI8GDwAA8J4Xy9LtceaRI0eaeMiQIc51p06divke999/f6HzOP/8853yoEGDTMySyvDVqlUr6hSyUsOGDZ3y888/b+Lp06c7dX/605/SktO59O7d2yn/8Y9/jHlttWrVTNyyZcuU5VRU2d+nDz30kFM3ZsyYAr9fXl6eUz527JiJy5UrF/N1gwcPdsphfEej4MaNG2fio0ePOnVPP/20iX/44Ye05ZQsengAAID3aPAAAADveTGkNXz4cBPbSyXDYnfBHjp0yKnLyckxcb9+/Zy6hx9+2MTBbl0khmGromf06NFO+corrzRxcNg4nUNapUqVcsq9evUy8fjx4506e4g5WDdhwgQTh7G1RKZp1KiRiZMZwhIR2bt3r4mHDRvm1O3Zs8fEc+fOTer9kT4VK1Y08YEDB5y60qVLpzmbwqGHBwAAeI8GDwAA8B4NHgAA4L2MnMPz05/+1Cl369Yt1PcPHi3xu9/9zsT/+Mc/nDp7PNqez4NwbN++/ayxiEizZs3SnU7Wsv9fX3/99TGvs49lSDd7SwqRM5c1x/Laa6855WSPlslUdevWdcozZ85M6HX23EZ73pOIyFNPPWXi3bt3O3X33XdfQu+/YMGChK5DuEqWLBmzbsuWLU75m2++SXU6oaKHBwAAeI8GDwAA8F7GDGnVqFHDxPPnz3fqateundB72F2rwVOAK1eubOIWLVo4dcme0tu+fXsTv/rqq0m9B360cuVKp8yQVvrYO+4Gl3/bkl3GnKwHHnjAxMHdlOPp3r27ideuXRtqTpnG3j5D5Mwhrljee+89Ez/66KMxrwvuQD9gwICY1y5btszEixYtSigPhOv22293ykr9ePD48uXL051OqOjhAQAA3qPBAwAAvEeDBwAAeK/IzuGpWrWqU37wwQdNfPHFFyf0Hv/617+ccqdOnUxctmxZp65Vq1YmDmupXePGjU3MHJ7Cs++fiMicOXMiysR/9onhIiJNmzY1sT2mL+L+3bbnYKRChQoVnPKoUaNMXL58eadu3759Ju7bt69TN2PGjBRkl5kKsgx/9erVJk709PLgcxvcVsRmL3sOHuOD9OjTp49TPn78uIlffPHFdKcTKnp4AACA92jwAAAA7xWpIS17uetzzz3n1NlLvOOxh7E6d+7s1Nm7IgdNnTo1ofcviHifh4JbsWJF1Clkjauuusop20Nc9knjIiIbN25MaS72MNbs2bNj1gXzsofaGMKKLTc31ynbw1ZBq1atMvFXX30V87q2bduaOHgavc0edhQRmT59esxrkR716tVzyvbvsUzfwoEeHgAA4D0aPAAAwHs0eAAAgPcincMTPJX18ccfN3Gic3ZE3DFo+5TeVJx6bB9BIXLmERW2Z599NvTPz2Y7duyIOoWsEW/7/8OHDzvl4LyasNnPd7yT2oNHzrz88sspy8knwfsZPD0+EcEtAZ544gkT5+TkxHzdb3/7W6f87rvvFvizUXilS5c2cfB32htvvJHudFKGHh4AAOA9GjwAAMB7kQ5pNWrUyCkPHjw45rV2t+s777zj1NnDWGHtkhxLv379nLLdlTt37lyn7uTJkynNJdvUrFkzbhnpkZeX55TDHjqeMGGCU7Z3fg0uPd+/f7+JJ02a5NQtWbIk1LzgsndMfumll5y64He7zd6N2z5xHdGxTwWoVKlSdImkGD08AADAezR4AACA99I+pNWwYUMTz5s3L+HXvf/++ybu2LFjmCmd080332zieKtX7IPvREROnTqVspyyUXAIi1VbqfPRRx855RYtWpg4eIDnmDFjTDxt2jSnLniAbyz2MFbv3r0TztNeIRY8lNIe/vruu++culTvDu0re6hq2LBhJr7yyitjvmblypVOuVevXiY+ePBgiNkhWfGGIL/++us0ZpJa9PAAAADv0eABAADeo8EDAAC8l/Y5PPbOmjVq1Ih5nT0vQERk8uTJKcvpXOw87RPdg7Zu3ZqGbLJXcC5As2bNIsrEfyNGjHDKN910k4mDpynbz3Rw51yllIntOTXBOTT23L6gYsV+/HdZcF5cz549TRzcMdh+Hu2df8/2+Ti7Cy64wCk//fTTJo73/NnP6vDhw526Tz75JKTsEJZ482LXrFmTxkxSix4eAADgPRo8AADAeykf0goud+vcuXPMa+0u6X/+859OXSoOAk3UfffdF7Puiy++MHFwt1GEi2Xp6RNcxn3HHXeYOLg7bnCZeiz2kNbll18esy7IHsYKXvf666+beP369U7dqFGjEsoLLnv4Mrh79aWXXnrW1+zZs8cp20ONDGEVfVWqVIlZZ+9mnuno4QEAAN6jwQMAALxHgwcAAHgv5XN4fvOb3zjlqlWrxrx26dKlJg7O4UmnQYMGOeWf//znMa+99dZbTXz06NGU5YQz5+x06tQpokyyjz0/pn79+k5d27ZtTWwfNyAi0qBBgwJ/1pEjR5zy0KFDTTxz5kynbu/evSbmKJfkNG3a1Cnb9zDWnB0R9/99ly5dnDrm7WSWeHNkfbqX9PAAAADv0eABAADeS8mQlr0bcffu3WNeF1zKeOedd6YinYTYJ6KPHj3aqStZsqSJN2/e7NRt27YttYnBCA5hzZkzJ6JMsps9lCEismDBAhO3bNnSqUt0SGvXrl0mDg6DL1u2rKAp4hzKli1r4nHjxjl1LVq0iPk6e7sCe6uCdExBsP8u/eIXv3Dq7J2z161bl/JcfGNPF8jLy3PqSpRI+4EMKUMPDwAA8B4NHgAA4D0aPAAAwHspGZy74oorTHzhhRfGvC441rpv375UpHNW9vbpIiJTpkwxcfBEdHveTps2bZy6gwcPpiA7nM2KFSuc8vbt2yPKBLa7777bxD169EjqPd5++20Tr127ttA5wZWTk+OUN2zYYOIaNWok/D7lypUz8V//+lcTT5s2zbnuyy+/LGiKIiLys5/9zMT2lh8i7sntO3fudOrsY0SYw1NwTZo0MfGmTZucOnt+XaajhwcAAHiPBg8AAPBeSoa0gl2RsdgnHafDwIEDTTx48GCnrlKlSiYOLj23h7G2bt2amuRwTsGdlmvVqhVRJtntoosucsrt2rUzsVLKqbPL9jBHcPfydA5nZ4vmzZubeMiQIU5dQYaxbPZwf+3atU1s74adDn379nXKr7zySlo/3wf2MGe1atVMvGjRoijSSQt6eAAAgPdo8AAAAO/R4AEAAN5LyRyeDz74IKHr4p1CXhDt27c3sT2vo2vXrs51jRs3NrF9XIRI/KXnzNspGoLL0AcMGBBRJtlt/PjxTvnGG280sdbaqdu/f7+J7eeUOTvhK168uFO2l2pfd9116U4ndL169TLxrFmzIszED/Xr1zdx5cqVTbx8+fIo0kkLengAAID3aPAAAADvpWRIa+HChQld161bN6d82223JfV59u6fwaGqWOwTnkXck9rZPbloYhl6dOrVq2fiDh06JPy63NxcE3/yySeh5gRX8FTrZIex7O9G+/6JJHcP7b87Iu4p6/FMnjzZKe/evdvEp06dKnAecNn/P23ly5dPcybpQw8PAADwHg0eAADgPRo8AADAeymZw3Py5EkTr1692qlr2rSpiYPb0AdP9C2s4HLysWPHmjh4uu/Ro0dD/WzAV8GTqu3lrUH2EQYVK1Y08ffffx92WlnP/t4VcZcX28dMiIi8+eabJh49erRTZ59WH8ZcmfXr18ctIxqNGjU668+//fbbNGeSPvTwAAAA79HgAQAA3kvJkNaJEydM3Lp1a6du8eLFJraHtwoiuJzuhRdeMPH06dNNvGXLlph5IfN07tw56hSylr0TeXCX20cffTTm61auXGnivLy88BODEfx+a9myZUSZIBPYUz7sHdG//vrrCLJJD3p4AACA92jwAAAA79HgAQAA3lPB042dSqViVyIttNbq3FclhvsZvbDuJ/cyejybfuHZ9Eese0kPDwAA8B4NHgAA4D0aPAAAwHs0eAAAgPdo8AAAAO/R4AEAAN6jwQMAALxHgwcAAHiPBg8AAPBe3J2WAQAAfEAPDwAA8B4NHgAA4D0aPAAAwHs0eAAAgPdo8AAAAO/R4AEAAN77f9lYXLh0Ll8YAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 720x288 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# iid\n",
    "label_dist = torch.zeros(10)\n",
    "for (x,y) in iid_client_train_loader[25]:\n",
    "    label_dist+= torch.sum(F.one_hot(y, num_classes=10), dim=0)\n",
    "print(\"iid: \", label_dist)\n",
    "view_10(x,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "fqZgr4FO41Bi"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "non-iid:  tensor([300.,   0.,   0., 300.,   0.,   0.,   0.,   0.,   0.,   0.])\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjwAAAD1CAYAAABUdy/PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAklElEQVR4nO3dd7hUxfnA8XdooQkKArEBFggo1UoICIqICoqCjZ89YCixRxJNFAUVBYkCgqgYI4JIExtFRUEUC3mwYMSCiCAiJkiV3s7vD5LhnQl77t69u3f3zn4/z5PneYd39+zI2VkmZ5qJokgAAABCVirbFQAAAMg0OjwAACB4dHgAAEDw6PAAAIDg0eEBAADBo8MDAACCR4cHAAAEL8gOjzFmnDFmlTFmozFmsTGmR7brhNQZY6oZY14wxmw2xiw3xvxftuuE1HAvw8JvbTjyoW2aEDceNMYcJyJLoijaboxpICJviUjHKIo+zG7NkApjzHOyt3PeXUSaich0EWkZRdGibNYLhce9DAu/teHIh7YZ5BOeKIoWRVG0/b/F//zv6CxWCSkyxlQSka4icmcURZuiKJonIi+LyBXZrRkKi3sZHn5rw5AvbTPIDo+IiDHmUWPMFhH5UkRWiciMLFcJqakvIruiKFqs/myhiByXpfogddzLAPFbG4S8aJvBdniiKOojIgeISGsRmSoi2+PfgRxVWUQ2en+2QfbeW5Qs3MsA8VsbhLxom8F2eEREoija/Z9Hc4eLSO9s1wcp2SQiVbw/qyIiP2ehLiga7mWg+K0t8fKibQbd4VHKCOPKJdViESljjKmn/qypiAQzkS6PcC/Dx29tyZQXbTO4Do8xpqYx5lJjTGVjTGljTAcR6SYib2a7bii8KIo2y97H5AOMMZWMMb8Rkc4iMja7NUNhcS/Dwm9tOPKlbQbX4ZG9qwR6i8j3IrJORIaIyE1RFL2c1VqhKPqISAUR+beIPCcivUNaKplnuJfh4Lc2LMG3zSD34QEAANBCfMIDAADgoMMDAACCR4cHAAAEjw4PAAAIHh0eAAAQvDJxSWMMS7iyLIoik65rcT+zL133k3uZfbTNsNA2w5HoXvKEBwAABI8ODwAACB4dHgAAEDw6PAAAIHh0eAAAQPDo8AAAgODR4QEAAMGjwwMAAIIXu/FgSXHuuefa+KWXXkr4OmP27UU0YcIEJ7d06VIbP/HEE05u+fLlRa0isqhNmzZO+a233rLxxx9/7OQ6duxo41WrVmW0Xii8vn37OuVzzjnHxhdeeKGTW7NmTbHUCUDJwBMeAAAQPDo8AAAgeHR4AABA8EwUJT7nrKQcgnbffffZ+Lbbbivy9ZYsWeKUO3ToYONly5YV+fqFwQGFqTn00ENt/Pnnnzu5Aw44wMb+93/mzJk21nPD0oUDCsNRkttmkyZNbKznQYm489j830Jtx44dNh4/fryTe++992y8c+fOlOtZnGib+1e9enUbz54928np71Hcb6n+ThUHDg8FAAB5iw4PAAAIXhDL0mfMmGHjdAxpHXPMMU75tddes/HJJ5/s5DZs2FDkz0PRderUySk3b97cxpUrV076OkcccYSN9aNcEZY5Z0vp0qVt7G8ZoYcrx4wZ4+R++umnzFasBPGX89999902Xrt2rZPTQxN16tRxcmXLlrVxrVq1bHzttdc6r/vuu+9sPGjQICen79OWLVsKqjqKWalS7nMQff8aNWrk5Pbs2ZPwOuXKlbOx/t6IZG+Ykyc8AAAgeHR4AABA8OjwAACA4AWxLF2PD3bu3NnG/lK4U045xca/+tWvUvqs448/3ikvXLgwpeskqyQvfc20Fi1a2Pjll192ctWqVUv4vpUrV9r4l7/8pZPT80UuvfRSJzdlypSU6qmx9HX/ypTZN53w9NNPd3J//vOfbewfE7Jp0yYb63lbIvFLqtOhJLXNYcOGOeVKlSrZuGfPnk5u9+7dCa+jt3TQ8xnbtm3rvE7PGdJzOUTcOZFdunRxclu3bk342ZlG29yrd+/eTnnEiBEJX/vjjz/auEqVKk6uYsWKNj7jjDOc3Jw5c4pSxQKxLB0AAOQtOjwAACB4QQxpJUs/cmvZsqWTO/HEE23cv3//hNd45JFHnPJNN92UnsolUJIemxe3iy66yMbPPfeck1u/fr2N7733Xic3dOhQGw8fPtzJ9enTx8YMaWVO7dq1nXK7du1s/NRTTyV9Hf1o3B8Ky7SS1Db9pcZa3NLiVNWoUcPG06ZNc3InnXSSjR966CEnd+utt6a9LsnK57b5zDPP2Pjyyy93cnoJ+dtvv+3krrnmGhufffbZTu6xxx6z8eTJk52c/9uabgxpAQCAvEWHBwAABI8ODwAACF4QR0ska+PGjTZ+9dVXnZx/nARyz8033+yU77zzThsb4w7Z9urVy8Zxc2/eeecdp/z73/++KFVEDH1MgV6aLBK/TcSuXbtsrJevixRuvk8+y8Q8nTirV6+2sd4qRETkhx9+sLF/JEw25/CEzt+CQ89H1ffI/67ccccdNn7wwQcTXv+LL75ImFuxYkXS9cwknvAAAIDg0eEBAADBy6shLU0vmxRxlyPH+fLLLzNRHSTQrFkzGw8ZMsTJ6ROZ/VO0p0+fntLnxW3TgMLT2z1MnDjRxkcddZTzOj1s5T82X7dunY0HDx7s5H7961/beNy4cUWrLDJC74ju+/bbb4uxJvnNH673d7n+L72cXCR+GEuLmxby7LPPJnWNTOMJDwAACB4dHgAAEDw6PAAAIHhBz+Fp1aqVU27durWN/RNhDzvssKSu+eKLLxa5XoinT2F+/vnnbayXs4qIjB071sb9+vXLfMVQIP9U5KlTp9pYn7S9fPly53V6nsDIkSOdXPfu3RN+XsOGDVOqJzJLHx0yYMCAhK/T7Rvpp4/cOf/885N6zxtvvJHSZ1188cUpva848YQHAAAEjw4PAAAIXhBDWm3btrXx+PHjbVy1alXndeXLl0/qekuWLHHKeims3kEU6aGHsEREJk2aZOMKFSrYWJ+oLSKycOHCIn/2gQce6JS3b99u482bNxf5+vlALz2fOXOmk9M7I+/evdvGgwYNcl43atSohNfXS1r17rAiInXr1rWxv9UEbbX4VKlSxSmPHj3axo0aNXJyb775po2ffvrpjNYr37Rp08Yp62GscuXKJXxf06ZNbfzZZ5+lvV65gic8AAAgeHR4AABA8OjwAACA4AUxh6d9+/Y2rlWrVkrX+Oabb/Z7PRH3CAOkhz4ywl+aquftnHnmmTZOx5wdnz4JWERk7dq1NvaXwWMvfeq5iHtkhH+auZ63o0/Cjpuz49u2bVvC3JFHHmljf67InDlzkv4MFN7ZZ59t47/97W9OTp/M7R/z0rNnTxvrI0WQGr2lyrBhw5xc3LwdfZzSokWL0l+xHMQTHgAAEDw6PAAAIHhBDGm9/fbbNr799tsTvs4YY2P/VOxSpfb1/Xbu3JnG2kFE5IQTTnDKs2bNsrG/fcAll1xi448++ijtddG7ZR9++OFO7r333rNxJobQQtC/f3+n7J98rnXo0MHGejlyJujvjQhDWqnSw5KdOnVycl27drVxt27dbKx/P0VEXnrpJRtff/31To6h4vTSw7qNGzdO+LpPP/3UKestXPx/D5Olhy7jfgdyBU94AABA8OjwAACA4AUxpKUfXb/yyis2Pv30053XVaxYMeE19GPBK664wsnpnZaRPL0SSw9hibg7s/orPPxVHemsh4jIueeea2N/Bd7NN9+c1s8OxX333WfjK6+8MuHr/vGPfzjlDz74IK312LJli1PWK/q2bt2a1s/KV3qoasyYMSld45hjjrHxcccd5+S+//771CoGEREpXbq0U9YHJ+tpGz6/Lf78889Frose0qpfv36Rr5dpPOEBAADBo8MDAACCR4cHAAAEL4g5PDt27LBx586dbdyqVSvndX379rVx3PwevROoiHui77///e8i1TVk/hiu/vv2l57rJcq33HKLk0vHXAx9Aru/k7Pmj2svWLCgyJ8dopYtW9rYnyewbNkyG3fs2NHJpfvE+ZUrVzrl6tWr2/jJJ59M62flq0mTJtn44IMPdnJ6Sbn+3a1du7bzOj3vcdq0aU5Oz5GcMGFC0Sqbh/SO5SIi7dq1s7G/vPzjjz+28W233Zb2uuhtJ1Jd2l6ceMIDAACCR4cHAAAEL4ghrUTmzZuXsHzttdc6uccee8zGdevWdXJ6CTVDWondc889Tlnvyvrss886uUcffdTGmzZtSunz9DDk3Xff7eQuvvhiGx9yyCFOTtdF1wP76GXFIiInnnhiwtfqg0B/+umntNdFD5U2aNDAya1atcrG69evT/tn56Pt27fb+OGHH07pGl9++aWNR48e7eQGDhxoY72zuQgHNSeiDwFt06ZN0u8bPny4jTds2JDWOomIHH/88Wm/ZibxhAcAAASPDg8AAAgeHR4AABC8jMzh0dvzt2/f3smtWLHCxn/84x+dXCbGGBOJmzfiz0Ngy/rE9HynJk2aJHxd3HEEqXr88cdtrLfDL0gm6hKa8uXLO+XKlSsnfO3q1aszWhc950rPZRBx54r4S9aRPa+99pqN/aNj7rrrLhsPGDDAyV199dUZrVdJpZf966XgBfFPSE+3uKMsNm7caONc+TeUJzwAACB4dHgAAEDw0jKkpXc7FXGHqmrWrJnwfaeccopT1jtBvvrqq+momuPss8+28SOPPJLwdfr0dREelcfJxGm5nTp1snG9evVs/NBDDzmv27NnT8Jr6GEQhrAK7/DDD0+Y27Vrl1P+4osvivx5eqjqoosucnL6O+Dj5O3cF7dVQYsWLYqxJki3uN2V9TYwX331VXFUp0A84QEAAMGjwwMAAIJHhwcAAAQvLXN42rZt65Tj5u1ojRs3dsp63sXy5cudnJ5zk+wJzOecc45TPu+882zsn96t6aWucB1xxBFOeejQoTaOG8/152Xoe1GnTh0n17x5cxtXqFDBxv6cHf15a9ascXIcGVE0W7ZsSZgrU8b92dDHePinzydywQUXOOUuXbrY+PLLL0/4vtmzZzvl6667LqnPg6tp06Y2bt26tZMbMWJEcVcHBdDt8ccff3Ryeh5lpvnHWpx22mk2XrdunZPLxe8RT3gAAEDw6PAAAIDgpWVI64033nDKS5YssbF/6nKcAw88cL+xiMiTTz6Z8H16t8e4YZU4esjs/vvvT+ka+cBfRhp3irY2YcIEp5zKffJPVv7kk09srE+7FxFZtGhRoa+PfT7//HOn/M0339j46KOPdnJxQ1Ca3nW9UaNGSddF/77oLQtE3JO9kTx96ry/bcPIkSNtnOrv6cEHH5zw+hrTB5Lzww8/2Pjdd991cl27dk34voYNG9pY/17G8dvmSSedZOPBgwc7uWrVqtl49OjRTk7vtp0reMIDAACCR4cHAAAEjw4PAAAIXlrm8PinnJ988sk21qdpi4hcc801NvbHHg899NB0VCchfarzkCFDnNzw4cNtvGPHjozWAy5/HsbXX39t43vvvdfGU6ZMKbY65Tv/OICpU6fa+IYbbnByNWrUsPHNN9+c0udNmzbNxo8//riTe+edd2zMnJ300NsH+EeD9O/f38b9+vVL6nr6OyAi8vrrr9tYL4EXcU/w7tatW1LXxz6FmVel5zb6/97+61//srGeX6fnX4nEb+Gij2HKxWXoPp7wAACA4NHhAQAAwTNxj8eMMamtSUyS3kVXRKRUqX39r969ezu5uMdqf/nLX2ys/3umT5/uvE4/nk12iV62RVFkCn5VctJxP/0T0Xv16mVjf6jD+2ynPGDAABv7S1MnTpxYlCrmtHTdz0y3zTj+rtmnnnqqjf2TzTt06GDjWbNm2VjffxGR999/38a7d+9OSz0zLdfaZqr8ISc93KWnAfj09iP+rvl6WGT+/PlOTu+OvWDBgsJVNoNKStv0t2zRQ/165+NU+b/Vc+fOtfHAgQOdnB5u3rZtW5E/O10S3Uue8AAAgODR4QEAAMGjwwMAAIKX1Tk8KFgo8wSwV0mZJ4CChdo2r776ahvr+VkiIpdddpmNy5Ytm/Aa+niKYcOGOTk99yeXlNS2qY9vuuqqq5xcjx49bFyzZk0np+dP6W0E9LwcEXfp+c6dO4tW2WLCHB4AAJC36PAAAIDgMaSV40J9bJ6vSupjc/wv2mZYaJvhYEgLAADkLTo8AAAgeHR4AABA8OjwAACA4NHhAQAAwaPDAwAAgkeHBwAABI8ODwAACB4dHgAAEDw6PAAAIHh0eAAAQPDo8AAAgODR4QEAAMGLPS0dAAAgBDzhAQAAwaPDAwAAgkeHBwAABI8ODwAACB4dHgAAEDw6PAAAIHh0eAAAQPCC7PAYY8YZY1YZYzYaYxYbY3pku05InTGmmjHmBWPMZmPMcmPM/2W7TkgNbTMstM1w5EPbDHLjQWPMcSKyJIqi7caYBiLyloh0jKLow+zWDKkwxjwnezvn3UWkmYhMF5GWURQtyma9UHi0zbDQNsORD20zyCc8URQtiqJo+3+L//nf0VmsElJkjKkkIl1F5M4oijZFUTRPRF4WkSuyWzOkgrYZDtpmWPKhbQbZ4RERMcY8aozZIiJfisgqEZmR5SohNfVFZFcURYvVny0UkeOyVB8UEW0zGLTNwITeNoPt8ERR1EdEDhCR1iIyVUS2x78DOaqyiGz0/myD7L23KIFom8GgbQYm9LYZbIdHRCSKot3/ecx6uIj0znZ9kJJNIlLF+7MqIvJzFuqCNKFtBoG2GaCQ22bQHR6ljAQ2FplHFotIGWNMPfVnTUWESZFhoG2WXLTNsAXXNoPr8BhjahpjLjXGVDbGlDbGdBCRbiLyZrbrhsKLomiz7H20OsAYU8kY8xsR6SwiY7NbMxQWbTMstM1w5EvbDG5ZujGmhohMkb3/T6OUiCwXkeFRFI3OasWQMmNMNRF5SkTai8gaEbktiqLx2a0VCou2GR7aZhjypW0G1+EBAADwBTekBQAA4KPDAwAAgkeHBwAABI8ODwAACB4dHgAAELwycUljDEu4siyKIpOua3E/sy9d95N7mX20zbDQNsOR6F7yhAcAAASPDg8AAAgeHR4AABA8OjwAACB4dHgAAEDw6PAAAIDg0eEBAADBo8MDAACCR4cHAAAEjw4PAAAIHh0eAAAQPDo8AAAgeLGHh2bTCSec4JRnzZpl46pVqyZ1jXfeeccpd+rUycabNm0qQu1QkkXRvrP99uzZ4+SWL19u4/vvv9/JjR49OrMVK6G6detm4+bNmyd8Xe/evW1cuXJlJ+ffh0S6d+/ulJ9++umk3gfAVb9+fRt36dLFxgMHDkzL9a+66iobjx07Ni3XLCqe8AAAgODR4QEAAMHL2SGtU0891SlXqVLFxnpIIk7r1q2d8hdffGHjpUuXOrnLLrvMxt9//33S9URuaNeunVO+8cYbbdy0aVMnp4dP/O9S7dq1bdyoUaN0VjFYZ511lo11O4rjD2El26affPLJhO8bM2ZMUtcA8oWeGvLII484ucaNG9u4YsWKNvbbojEmYS7O7bffbuPvvvvOyc2dOzfp66QTT3gAAEDw6PAAAIDg0eEBAADBy9k5PJlwyCGH7DcWcZe9+0tdBw0alNF6ITmHHnqoU9bzOVq2bOnk/GXPyJwHH3zQxuPGjUv4ul69etl4wYIFTs4vJ/Lcc885Zb11AHN4ckf16tVtXLZs2YSvW716tVPevXt3xuqUj4499lgbt2jRwsklmo+zbt06p6y36njjjTcSftaFF17olBs0aGDjVq1aOTnm8AAAAGQIHR4AABA8E7fMzBiT/Bq0NNO7QIqITJ482cZ6ibqI+/hUL6/Ty+lECrekTrvppptsPGrUKCeX6UewURSZgl+VnGzez3SYM2eOU/a3HUhWssssR44c6ZT1UvdUpet+lvR7maoffvghYc4f8sw02uY+Z5xxhlOeMmWKjQ844AAnp9vfCy+84OT0Lr8ffvhhOqtYoBDb5m9+8xsb+0O+a9eutfHzzz9vY3/rhzVr1iT1Wf6Q1qRJk2ysh8VERI4//ngb+0No6ZDoXvKEBwAABI8ODwAACB4dHgAAELycncNTGB07drRxvXr1bOwfT3HeeeeldH095qyPHhARWblyZUrXTBbzBPZhDs8+Jf1eFsabb75pY/+ev/rqqzZOtX2nira5z3vvveeU/SXQydq6dauN9XYHIiJ33313StdMFm0zvfRvqz+Hp3nz5jZmDg8AAEAa0eEBAADBC2Kn5enTp9u4TZs2Nm7btq3zOn+ZerJKlaJfWJw6depkYz1EuWHDBud1cfdTfycWL17s5E488UQbxw2LHXTQQQVXFmnXrFkzp6xPrZ83b56TK+5hrHxTs2ZNp6x3qNdLi0866STndQVMlUj4uvLly9v4zjvvdHKVKlWy8fDhw53cihUrEn4eiseECROcsr63o0ePdnKZGMZKBv+SAwCA4NHhAQAAwaPDAwAAgldi5vDMnz/fxtu2bXNyem6HXorub2me6tESe/bsSel9SEzP02nYsKGT69evn40rVKhg47ijQmbPnu3k9Bb1H3zwgZPTcw9mzZrl5KpWrbrf6yOzTjjhBBu/+OKLTk7PodPfDaROz0+74447nNzpp59u44MPPtjJ6eM7kt3ewdelSxcbL1261MndddddNj7//POd3B/+8AcbX3rppU7u+uuvt7H//UHB9Pwo/99NPd/Gn9fYt29fG5911llOTv+7PHHixLTUs6h4wgMAAIJHhwcAAAQvZ4e0brvtNqfcuHFjG5crV65Y66JPaN6xY0exfnYo9HYBIiLjxo2zceXKlZO6hn9Sdvfu3W3s7/S6adOmhNdZtWqVjYcMGeLk7rnnHhsfeOCBTk7/N8ydO7fgCiMhf+m5HobQS59FRC644AIb+8vSkZgeftBDDyIiPXr0sLE/bJWOodxFixbZ2G9jcUNO77//vo39IS1ND62JuEvYGdLaR59ir4eNRUROO+00G9epU8fG9evXd1738ccf21jvkFwQPb1kyZIlSb8vk3jCAwAAgkeHBwAABI8ODwAACF7OzuHxx5yLe96ONmPGDBuvXbs2a/XIddWrV3fKFStWtHGfPn2cXOnSpQt9fX978tdff73Q1xBx5/Dok7hF3LkNPubtFM1FF11k4xEjRjg5/7ujDR482Mb+URLXXnttmmoXHn38wmWXXZbwdc8884xT1m1i2rRpTk4vUdbvu/zyy53X6VPs/evHGTlypI31UmkRd8m6v0WFnlvSvn17J+dvPREyPS9HRGTMmDE29ufGJUtv4xHH32Lg3XffTenzMoknPAAAIHh0eAAAQPBMAafaZm2rWf8xmn4sqXfDjRO3M29h6OvUrl3bya1cuTKlayYriqLUjnjfj0zfT/+xuX6cmo57UaZM+kdg/VOX9dDbzz//7OT0MtlUh7fSdT+z2TYL44EHHrDxLbfcYmN/SDPVtvn555/buEmTJildI1W51jb18n0RkUmTJtlY71YtIvLhhx/a+OSTT07p84YOHWpjvdOxiLsU/U9/+lNK1/fprQz89qe3tvjoo4+cnH+SeyIhtM05c+Y4Zb00PO43eOfOnTbWbUpE5Nlnn7Wxv1XH1VdfbWN/qwC93UFxn46e6F7yhAcAAASPDg8AAAhezg5pFbfrrrvOxsOGDXNy+nGwfwiaf4hduuXaY3Nf3bp1baxXZoiI1KtXz8b+I/W4A1n1So0bbrihiDWM99133znlww47zMbr1693cnr1h//YPFkhPDaPc+aZZzrlmTNn7vd1mzdvdsrffvutjceOHZvw+npHXRH3e+Xv5p3qPUpWrrVNPYQlItK1a1cbr1mzxsnpXXdXrFiR0ufpQyb1zsoi7m7HmWjD/q7rtWrVsvGWLVucnH8YZiIhtM0aNWo4ZX/1nKZX4+3atcvG/pBWnAYNGtjY/w7o4a64Np0JDGkBAIC8RYcHAAAEjw4PAAAIXs7utFzcRo0aZeNGjRo5Ob2ba8eOHZ2cXj6f6TkDuUgvTdXzX0TcZY/+nB2de+qpp5zcwIED01jD//XKK6/YWI/9+/UaP368k1u8eHFG6xWCjRs3OuWFCxfa+LHHHrOxf+p5svMG9NJ2EXeH5lR3kg3FhRde6JT1d1kv2xZx53qkOodHb9ugd9EWEenevbuN/fuidzovDL3suUKFCglfp0/3zjerV692yg8//HBGP0/PvfP//TvqqKMy+tmp4AkPAAAIHh0eAAAQvKwOaf3iF79wyvoAOJ9+XLZjx46012X37t023rRpU8LX+Y9Sy5Ytm/a65LJ27do55bZt29o47jGzT+9a/NZbbzm5uL//ZOmhqqlTpzo5/Yg97hBT/zC8dNQrdB988IFTTvbgwTh6KMM/RFgvt54+fXqRPyskemfd8uXLO7m4w1pTMX/+/NhyKvxhMr1rt7/b/rJly2zs756OzNm+fbuN/YO1W7ZsWdzVKRBPeAAAQPDo8AAAgODR4QEAAMEr9jk8et7OHXfc4eRuv/12G/snu06bNs3G99xzj5NbsGBBOquIGDfeeKNT9pe7JtKzZ0+nrLc137p1a5Hr5Z/i26NHDxufcsopTk5/t/yjVR588EEbT548ucj1QuH5py6/9NJLNvbvsz//C/vEHRuUzSXD+qiHvn37Ojk9j/OMM85wcvrfjmeeecbJ6SNHUl1mn0t69+7tlMeNG2djvR1ALon7vuUKnvAAAIDg0eEBAADBK/YhLf3IUg9hFUTvcFytWjUn17p16yLXSz8q909dznf67+PUU09N+n16KOLrr792cukYxmrWrJmNR48e7eSSHWr75JNPEpa///77VKuGQqpZs6aN9U7YIiJNmzZN+L7OnTtnrE4lzeDBg52yP1ykPfroozZu0qSJk9O75+qpBCIi69ev3+/1/KHGTp062fjII490cmeddZaN69at6+T0juz+Fgd//etfE9ZLn/ZdUtWvX9/GQ4YMcXIXXHCBjfv16+fk/vnPf9p48+bNGard/untDVq0aOHk/B30cwFPeAAAQPDo8AAAgODR4QEAAMEr9jk8Dz30kI39pedaqVJuX0yP7fpbVv/973+38aeffmpj//RWPUZ63HHHObnrr78+qbr4p37H/TeEQh+pEPff72/tr7d4nzt3btKfp8f/69WrZ2N/btW5556b9DU1fZqyv6zZnzeAzLjkkkucst6i4thjj3Vy27Zts7F/Inhxz1nIZXpLBRH3GJi4Iz78JdB6efGgQYMSvi9ue4c4+n3+cQT33XefjTN90neu0Ufz+EeB6HvpL9fXv1l6rquIyLp169JZxf+h543p7QZERJ5//vmMfnYqeMIDAACCR4cHAAAEr9iHtAYMGGBjf2mh5g+dxD0yveKKK/b75/7jbv2Y0D8lO+76ui7+kll/2CxEX331lY39v9MqVarYuE6dOk4ubhhLD1s1bNjQyelll/oxrz98mOrOnnqozd+xFXsddNBBTrlLly429ocWH3/8cRvrIWURd0ffXr162fh3v/ud8zp9L/37+tvf/tbGM2fOLLDu+UqfHC8icuaZZ9rYHwrs2rWrjTt06FDkz/bb5vvvv29jvT2FiMiUKVNsvHTp0iJ/dij0Dsp6GoGIO1zk/12nY1qF3sXab996ewO/berhNX9o9N133y1yvdKNJzwAACB4dHgAAEDw6PAAAIDgmbh5EMaYtB9/2qhRIxv7W0/rpZPpmK+Rrjkfq1atsrHe4lsk8ye1R1GUtnXv6bifev6LiEifPn1svGPHDienl3/790IfC1GuXLmkPjvufvonCM+bN8/G9957r5ObP39+Up+XCem6n5lom9qLL77olPWcK/8+fPbZZzb2vwP6KBnNv4a+f7feequTmzhx4n5fl2251jZTVatWrbRfU88nKinHPuRS22zfvr1Tnjx5so31vEkRt01MmjTJyell6bNnz3Zyeh6X3uKjRo0azut0W/XncD7wwAM29rdF8H8LilOie8kTHgAAEDw6PAAAIHjFPqSlVaxY0SmPGjXKxvoRuohI1apVC339VIe0/BOB9S6XCxcuLHQ9iiLXHpv7y1v1Tqz+snTvs51yKsOLK1eudMp6u4D+/fs7uaeffrrQ1y8OufTYPI4ejhQRady4sf5sJ5fKvXziiSecst6BfcmSJYW+XjbkWttE0eRy29TbO/jL/P2dyROJa7f6t3T16tXO68aNG7ffWOR/t6HIFQxpAQCAvEWHBwAABI8ODwAACF5W5/DE8U/3ve6662x85ZVXJnWNuDHLZcuWObkRI0bY+O2333Zy2Tw+ItfmCegT50VEzjnnHBv7R3w0bdpUf7aTi/ve6TlUekn50KFDC1PVnJTL8wS0Vq1aOeUZM2bYuFKlSk5O38vly5c7OX2MgJ6j57e/kijX2iaKpqS0zZ49ezplPf/NP2Vd83+D9ZEwEyZMsHHckUAlBXN4AABA3qLDAwAAgpezQ1rYi8fmYSkpj81RMNpmWGib4WBICwAA5C06PAAAIHh0eAAAQPDo8AAAgODR4QEAAMGjwwMAAIJHhwcAAASPDg8AAAgeHR4AABA8OjwAACB4dHgAAEDw6PAAAIDg0eEBAADBiz0tHQAAIAQ84QEAAMGjwwMAAIJHhwcAAASPDg8AAAgeHR4AABA8OjwAACB4/w/lECnaRuUyoQAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 720x288 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# non-iid\n",
    "label_dist = torch.zeros(10)\n",
    "for (x,y) in noniid_client_train_loader[25]:\n",
    "    label_dist+= torch.sum(F.one_hot(y,num_classes=10), dim=0)\n",
    "print(\"non-iid: \", label_dist)\n",
    "view_10(x,y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2cAwYqdeHvcY"
   },
   "source": [
    "## Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pnKxxiKYo4cd"
   },
   "outputs": [],
   "source": [
    "# define fully connected NN\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MLP, self).__init__()\n",
    "        self.fc1 = nn.Linear(784, 200);\n",
    "        self.fc2 = nn.Linear(200, 200);\n",
    "        self.out = nn.Linear(200, 10);\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.flatten(1) # [B x 784]\n",
    "        x = F.relu(self.fc1(x)) # [B x 200]\n",
    "        x = F.relu(self.fc2(x)) # [B x 200]\n",
    "        x = self.out(x) # [B x 10]\n",
    "        return x\n",
    "\n",
    "print(MLP())\n",
    "print(num_params(MLP()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jefi6wlYpXYa"
   },
   "outputs": [],
   "source": [
    "# define cnn\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 5) \n",
    "        self.conv2 = nn.Conv2d(32, 64, 5)\n",
    "        self.fc = nn.Linear(1024, 512)\n",
    "        self.out = nn.Linear(512, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.max_pool2d(self.conv1(x), 2, 2) # [B x 32 x 12 x 12]\n",
    "        x = F.max_pool2d(self.conv2(x), 2, 2) # [B x 64 x 4 x 4]\n",
    "        x = x.flatten(1) # [B x 1024]\n",
    "        x = F.relu(self.fc(x)) # [B x 512]\n",
    "        x = self.out(x) # [B x 10]\n",
    "        return x\n",
    "\n",
    "print(CNN())\n",
    "print(num_params(CNN()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r6Tmim2lHz4j"
   },
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A4mTFWgx0gmq"
   },
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "def validate(model):\n",
    "    model = model.to(cuda)\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for (t, (x,y)) in enumerate(test_loader):\n",
    "            x = x.to(cuda)\n",
    "            y = y.to(cuda)\n",
    "            out = model(x)\n",
    "            correct += torch.sum(torch.argmax(out, dim=1) == y).item()\n",
    "            total += x.shape[0]\n",
    "    return correct/total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-YDOK1S0H_Xe"
   },
   "outputs": [],
   "source": [
    "def train_client(id, client_loader, global_model, num_local_epochs, lr):\n",
    "    local_model = copy.deepcopy(global_model)\n",
    "    local_model = local_model.to(cuda)\n",
    "    local_model.train()\n",
    "    optimizer = torch.optim.SGD(local_model.parameters(), lr=lr)\n",
    "\n",
    "    for epoch in range(num_local_epochs):\n",
    "        for (i, (x,y)) in enumerate(client_loader):\n",
    "            x = x.to(cuda)\n",
    "            y = y.to(cuda)\n",
    "            optimizer.zero_grad()\n",
    "            out = local_model(x)\n",
    "            loss = criterion(out, y)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "    return local_model\n",
    "\n",
    "def running_model_avg(current, next, scale):\n",
    "    if current == None:\n",
    "        current = next\n",
    "        for key in current:\n",
    "            current[key] = current[key] * scale\n",
    "    else:\n",
    "        for key in current:\n",
    "            current[key] = current[key] + (next[key] * scale)\n",
    "    return current"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pJgWgjM-Pv4R"
   },
   "outputs": [],
   "source": [
    "def fed_avg_experiment(global_model, num_clients_per_round, num_local_epochs, lr, client_train_loader, max_rounds, filename):\n",
    "    round_accuracy = []\n",
    "    for t in range(max_rounds):\n",
    "        print(\"starting round {}\".format(t))\n",
    "\n",
    "        # choose clients\n",
    "        clients = np.random.choice(np.arange(100), num_clients_per_round, replace = False)\n",
    "        print(\"clients: \", clients)\n",
    "\n",
    "        global_model.eval()\n",
    "        global_model = global_model.to(cpu)\n",
    "        running_avg = None\n",
    "\n",
    "        for i,c in enumerate(clients):\n",
    "            # train local client\n",
    "            print(\"round {}, starting client {}/{}, id: {}\".format(t, i+1,num_clients_per_round, c))\n",
    "            local_model = train_client(c, client_train_loader[c], global_model, num_local_epochs, lr)\n",
    "\n",
    "            # add local model parameters to running average\n",
    "            running_avg = running_model_avg(running_avg, local_model.state_dict(), 1/num_clients_per_round)\n",
    "        \n",
    "        # set global model parameters for the next step\n",
    "        global_model.load_state_dict(running_avg)\n",
    "\n",
    "        # validate\n",
    "        val_acc = validate(global_model)\n",
    "        print(\"round {}, validation acc: {}\".format(t, val_acc))\n",
    "        round_accuracy.append(val_acc)\n",
    "\n",
    "        if (t % 10 == 0):\n",
    "          np.save(filename+'_{}'.format(t)+'.npy', np.array(round_accuracy))\n",
    "\n",
    "    return np.array(round_accuracy)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SrmY-TPnGI77"
   },
   "source": [
    " ## MLP experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Mwymk72xlaDP"
   },
   "outputs": [],
   "source": [
    "mlp = MLP()\n",
    "print(mlp)\n",
    "print(\"total params: \", num_params(mlp))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SDjaFBJAjeHC"
   },
   "outputs": [],
   "source": [
    "# MLP - iid - m=10 experiment\n",
    "mlp_iid_m10 = copy.deepcopy(mlp)\n",
    "acc_mlp_iid_m10 = fed_avg_experiment(mlp_iid_m10, num_clients_per_round=10, \n",
    "                                 num_local_epochs=1,\n",
    "                                 lr=0.05,\n",
    "                                 client_train_loader = iid_client_train_loader,\n",
    "                                 max_rounds=100,\n",
    "                                 filename='./acc_mlp_iid_m10')\n",
    "print(acc_mlp_iid_m10)\n",
    "np.save('./acc_mlp_iid_m10.npy', acc_mlp_iid_m10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u5x4lhIcyyt3"
   },
   "outputs": [],
   "source": [
    "# MLP - iid - m=50 experiment\n",
    "mlp_iid_m50 = copy.deepcopy(mlp)\n",
    "acc_mlp_iid_m50 = fed_avg_experiment(mlp_iid_m50, num_clients_per_round=50, \n",
    "                                 num_local_epochs=1,\n",
    "                                 lr=0.05,\n",
    "                                 client_train_loader = iid_client_train_loader,\n",
    "                                 max_rounds=100,\n",
    "                                 filename='./acc_mlp_iid_m50')\n",
    "print(acc_mlp_iid_m50)\n",
    "np.save('./acc_mlp_iid_m50.npy', acc_mlp_iid_m50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O4NeL5U4zKDS"
   },
   "outputs": [],
   "source": [
    "# MLP - non-iid - m=10 experiment\n",
    "mlp_noniid_m10 = copy.deepcopy(mlp)\n",
    "acc_mlp_noniid_m10 = fed_avg_experiment(mlp_noniid_m10, num_clients_per_round=10, \n",
    "                                 num_local_epochs=1,\n",
    "                                 lr=0.05,\n",
    "                                 client_train_loader = noniid_client_train_loader,\n",
    "                                 max_rounds=300,\n",
    "                                 filename = './acc_mlp_noniid_m10')\n",
    "print(acc_mlp_noniid_m10)\n",
    "np.save('./acc_mlp_noniid_m10.npy', acc_mlp_noniid_m10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O0MeL-jQzXd7"
   },
   "outputs": [],
   "source": [
    "# MLP - noniid - m=50 experiment\n",
    "mlp_noniid_m50 = copy.deepcopy(mlp)\n",
    "acc_mlp_noniid_m50 = fed_avg_experiment(mlp_noniid_m50, num_clients_per_round=50, \n",
    "                                 num_local_epochs=1,\n",
    "                                 lr=0.05,\n",
    "                                 client_train_loader = noniid_client_train_loader,\n",
    "                                 max_rounds=300,\n",
    "                                 filename='./acc_mlp_noniid_m50')\n",
    "print(acc_mlp_noniid_m50)\n",
    "np.save('./acc_mlp_noniid_m50.npy', acc_mlp_noniid_m50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yR4UdghyGNCN"
   },
   "source": [
    "## CNN Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RQKqbwx6413q"
   },
   "outputs": [],
   "source": [
    "cnn = CNN()\n",
    "print(cnn)\n",
    "print(\"total params: \", num_params(cnn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JEuqsiF3GkXC"
   },
   "outputs": [],
   "source": [
    "# CNN - iid - m=10 experiment\n",
    "cnn_iid_m10 = copy.deepcopy(cnn)\n",
    "acc_cnn_iid_m10 = fed_avg_experiment(cnn_iid_m10, num_clients_per_round=10, \n",
    "                                 num_local_epochs=5,\n",
    "                                 lr=0.01,\n",
    "                                 client_train_loader = iid_client_train_loader,\n",
    "                                 max_rounds=100,\n",
    "                                 filename='./acc_cnn_iid_m10')\n",
    "print(acc_cnn_iid_m10)\n",
    "np.save('./acc_cnn_iid_m10.npy', acc_cnn_iid_m10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-0Eo_EYbGTfY"
   },
   "outputs": [],
   "source": [
    "# CNN - iid - m=50 experiment\n",
    "cnn_iid_m50 = copy.deepcopy(cnn)\n",
    "acc_cnn_iid_m50 = fed_avg_experiment(cnn_iid_m50, num_clients_per_round=50, \n",
    "                                 num_local_epochs=5,\n",
    "                                 lr=0.01,\n",
    "                                 client_train_loader = iid_client_train_loader,\n",
    "                                 max_rounds=100,\n",
    "                                 filename='./acc_cnn_iid_m50')\n",
    "print(acc_cnn_iid_m50)\n",
    "np.save('./acc_cnn_iid_m50.npy', acc_cnn_iid_m50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-EUw04QbHN9s"
   },
   "outputs": [],
   "source": [
    "# CNN - non-iid - m=10 experiment\n",
    "cnn_noniid_m10 = copy.deepcopy(cnn)\n",
    "acc_cnn_noniid_m10 = fed_avg_experiment(cnn_noniid_m10, num_clients_per_round=10, \n",
    "                                 num_local_epochs=5,\n",
    "                                 lr=0.01,\n",
    "                                 client_train_loader = noniid_client_train_loader,\n",
    "                                 max_rounds=200,\n",
    "                                 filename='./acc_cnn_noniid_m10')\n",
    "print(acc_cnn_noniid_m10)\n",
    "np.save('./acc_cnn_noniid_m10.npy', acc_cnn_noniid_m10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gICkrXpMHdZT"
   },
   "outputs": [],
   "source": [
    "# CNN - non-iid - m=50 experiment\n",
    "cnn_noniid_m50 = copy.deepcopy(cnn)\n",
    "acc_cnn_noniid_m50 = fed_avg_experiment(cnn_noniid_m50, num_clients_per_round=50, \n",
    "                                 num_local_epochs=5,\n",
    "                                 lr=0.01,\n",
    "                                 client_train_loader = noniid_client_train_loader,\n",
    "                                 max_rounds=100,\n",
    "                                 filename='./acc_cnn_noniid_m50')\n",
    "print(acc_cnn_noniid_m50)\n",
    "np.save('./acc_cnn_noniid_m50.npy', acc_cnn_noniid_m50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BZaLNsaS00jV"
   },
   "outputs": [],
   "source": [
    "# view_10(x_debug[:10].to(cpu), torch.argmax(model(x_debug),dim=1)[:10].to(cpu))\n",
    "\n",
    "# m = CNN().to(cuda)\n",
    "# m.train()\n",
    "# lr = 0.01\n",
    "# opt = torch.optim.SGD(m.parameters(), lr)\n",
    "\n",
    "# for epoch in range(5):\n",
    "#     for (t, (x,y)) in enumerate(train_loader):\n",
    "#         x = x.to(cuda)\n",
    "#         y = y.to(cuda)\n",
    "#         opt.zero_grad()\n",
    "#         out = m(x)\n",
    "#         loss = criterion(out, y)\n",
    "#         loss.backward()\n",
    "#         opt.step()\n",
    "\n",
    "#         if (t%100 == 0):\n",
    "#             print(\"epoch {}, step {}, loss: {}\".format(epoch, t, loss))\n",
    "\n",
    "#     print(\"running validation\")\n",
    "#     acc = validate(m)\n",
    "#     print(\"epoch {} validation acc: {}\".format(epoch, acc))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "FedAvg",
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.9",
   "language": "python",
   "name": "base"
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}