baseline_mnist_with_clients.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision.transforms import ToTensor\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import DataLoader, TensorDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_augment_aggregate(loader, num_of_clients):\n",
    "    data, labels = next(iter(loader))\n",
    "\n",
    "    partite_class_data = []\n",
    "    partite_class_labels = []\n",
    "    for class_idx in range(10):\n",
    "        mask = labels == class_idx\n",
    "        len_mask = torch.count_nonzero(mask)\n",
    "        remainder = len_mask % num_of_clients\n",
    "        split_param = len_mask // num_of_clients\n",
    "        partite_class_data.append(list(torch.split(data[mask][:-remainder], split_param)))\n",
    "        partite_class_labels.append(list(torch.split(labels[mask][:-remainder], split_param)))\n",
    "        for remainder_idx in range(remainder):\n",
    "            partite_class_data[-1][remainder_idx] = torch.concatenate([partite_class_data[-1][remainder_idx], data[mask][-remainder].unsqueeze(0)])\n",
    "            partite_class_labels[-1][remainder_idx] = torch.concatenate([partite_class_labels[-1][remainder_idx], labels[mask][-remainder].unsqueeze(0)])\n",
    "\n",
    "    partite_clients_data = []\n",
    "    partite_clients_labels = []\n",
    "    for client_idx in range(num_of_clients):\n",
    "        client_data_buffer = []\n",
    "        client_labels_buffer = []\n",
    "        for class_idx in range(10):\n",
    "            client_data_buffer.append(partite_class_data[class_idx][client_idx])\n",
    "            client_labels_buffer.append(partite_class_labels[class_idx][client_idx])\n",
    "        client_data_buffer = torch.concatenate(client_data_buffer)\n",
    "        client_labels_buffer = torch.concatenate(client_labels_buffer)\n",
    "        \n",
    "        #normalize client data\n",
    "        client_data_buffer = (client_data_buffer - torch.mean(client_data_buffer)) / torch.std(client_data_buffer)\n",
    "        partite_clients_data.append(client_data_buffer.reshape(client_data_buffer.shape[0], -1))\n",
    "        partite_clients_labels.append(client_labels_buffer)\n",
    "    permute_data = torch.randperm(data.shape[0])\n",
    "    data = torch.concatenate(partite_clients_data)[permute_data]\n",
    "    labels = torch.concatenate(partite_clients_labels)[permute_data]\n",
    "    dataset = TensorDataset(data, labels)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_of_clients = 64\n",
    "transform = ToTensor()\n",
    "\n",
    "train_dataset = MNIST('./data', train=True, transform=transform, download=True)\n",
    "train_loader = DataLoader(train_dataset, batch_size=train_dataset.data.shape[0])\n",
    "train_dataset = collect_augment_aggregate(train_loader, num_of_clients)\n",
    "train_loader = DataLoader(train_dataset, batch_size=len(train_dataset))\n",
    "\n",
    "test_dataset = MNIST('./data', train=False, transform=transform, download=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=test_dataset.data.shape[0])\n",
    "test_dataset = collect_augment_aggregate(test_loader, num_of_clients)\n",
    "test_loader = DataLoader(test_dataset, batch_size=len(test_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleNetwork(nn.Module):\n",
    "    def __init__(self, in_channel, hidden_channel, out_channel, *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.lin_01 = nn.Linear(in_channel, hidden_channel, bias=False)\n",
    "        self.lin_02 = nn.Linear(hidden_channel, out_channel, bias=False)\n",
    "        self.act = nn.ReLU()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        out = self.act(self.lin_01(x))\n",
    "        out = self.lin_02(out)\n",
    "        return out\n",
    "    \n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model = SimpleNetwork(784, 128, 10).to(device)\n",
    "optim = Adam(model.parameters(), lr=0.001)\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.337559938430786\n",
      "epoch: 1, accuracy: 0.41429999470710754\n",
      "2.0882413387298584\n",
      "1.8813409805297852\n",
      "1.6887352466583252\n",
      "1.5054283142089844\n",
      "1.3369925022125244\n",
      "1.1879539489746094\n",
      "1.0588058233261108\n",
      "0.9487942457199097\n",
      "0.8575913310050964\n",
      "0.7834738492965698\n",
      "epoch: 11, accuracy: 0.8463000059127808\n",
      "0.722058117389679\n",
      "0.6693609356880188\n",
      "0.6244329810142517\n",
      "0.5872538685798645\n",
      "0.5566978454589844\n",
      "0.5309362411499023\n",
      "0.5084362030029297\n",
      "0.4884193539619446\n",
      "0.4707747995853424\n",
      "0.4555896818637848\n",
      "epoch: 21, accuracy: 0.88919997215271\n",
      "0.44261205196380615\n",
      "0.4310884475708008\n",
      "0.4203794002532959\n",
      "0.4104326069355011\n",
      "0.4014269709587097\n",
      "0.3933083713054657\n",
      "0.3858363926410675\n",
      "0.37886616587638855\n",
      "0.3723679482936859\n",
      "0.3662857115268707\n",
      "epoch: 31, accuracy: 0.9060999751091003\n",
      "0.3605315089225769\n",
      "0.35508760809898376\n",
      "0.34997889399528503\n",
      "0.34518036246299744\n",
      "0.34063535928726196\n",
      "0.3363131582736969\n",
      "0.3321903347969055\n",
      "0.32822850346565247\n",
      "0.3244074583053589\n",
      "0.3207532465457916\n",
      "epoch: 41, accuracy: 0.9172999858856201\n",
      "0.3172879219055176\n",
      "0.31398671865463257\n",
      "0.31080543994903564\n",
      "0.3077155351638794\n",
      "0.3047127425670624\n",
      "0.3018101453781128\n",
      "0.29902344942092896\n",
      "0.2963433563709259\n",
      "0.29373878240585327\n",
      "0.2911805212497711\n",
      "epoch: 51, accuracy: 0.9232999682426453\n",
      "0.2886695861816406\n",
      "0.2862243056297302\n",
      "0.2838556170463562\n",
      "0.281546026468277\n",
      "0.27927759289741516\n",
      "0.2770472466945648\n",
      "0.2748624384403229\n",
      "0.27272069454193115\n",
      "0.2706174850463867\n",
      "0.26854702830314636\n",
      "epoch: 61, accuracy: 0.9312999844551086\n",
      "0.26650795340538025\n",
      "0.2644994258880615\n",
      "0.2625194787979126\n",
      "0.26056715846061707\n",
      "0.2586384117603302\n",
      "0.2567325830459595\n",
      "0.2548496127128601\n",
      "0.2529911696910858\n",
      "0.2511543035507202\n",
      "0.24933528900146484\n",
      "epoch: 71, accuracy: 0.9346999526023865\n",
      "0.24753205478191376\n",
      "0.2457427829504013\n",
      "0.24396716058254242\n",
      "0.24220676720142365\n",
      "0.2404624968767166\n",
      "0.2387327253818512\n",
      "0.2370145469903946\n",
      "0.23530592024326324\n",
      "0.23360958695411682\n",
      "0.23192523419857025\n",
      "epoch: 81, accuracy: 0.9386999607086182\n",
      "0.23025524616241455\n",
      "0.2285986840724945\n",
      "0.2269544005393982\n",
      "0.22532325983047485\n",
      "0.22370408475399017\n",
      "0.2220965027809143\n",
      "0.22049981355667114\n",
      "0.21891440451145172\n",
      "0.21734122931957245\n",
      "0.21577996015548706\n",
      "epoch: 91, accuracy: 0.9422999620437622\n",
      "0.21423131227493286\n",
      "0.21269619464874268\n",
      "0.21117283403873444\n",
      "0.20966364443302155\n",
      "0.20816776156425476\n",
      "0.20668509602546692\n",
      "0.20521585643291473\n",
      "0.20375972986221313\n",
      "0.2023170292377472\n",
      "0.2008855640888214\n",
      "epoch: 101, accuracy: 0.9449999928474426\n",
      "0.1994662582874298\n",
      "0.19806012511253357\n",
      "0.19666652381420135\n",
      "0.19528670608997345\n",
      "0.19392067193984985\n",
      "0.19256621599197388\n",
      "0.1912255734205246\n",
      "0.18989887833595276\n",
      "0.18858769536018372\n",
      "0.18728958070278168\n",
      "epoch: 111, accuracy: 0.9479999542236328\n",
      "0.18600481748580933\n",
      "0.18473494052886963\n",
      "0.18347734212875366\n",
      "0.18223296105861664\n",
      "0.1810024082660675\n",
      "0.1797858327627182\n",
      "0.1785823404788971\n",
      "0.17739278078079224\n",
      "0.17621636390686035\n",
      "0.17505334317684174\n",
      "epoch: 121, accuracy: 0.9510999917984009\n",
      "0.17390234768390656\n",
      "0.1727638989686966\n",
      "0.17163649201393127\n",
      "0.17052152752876282\n",
      "0.1694180965423584\n",
      "0.1683260202407837\n",
      "0.1672447919845581\n",
      "0.1661745309829712\n",
      "0.16511446237564087\n",
      "0.16406501829624176\n",
      "epoch: 131, accuracy: 0.9529999494552612\n",
      "0.16302600502967834\n",
      "0.16199813783168793\n",
      "0.16098082065582275\n",
      "0.1599736362695694\n",
      "0.15897612273693085\n",
      "0.15798813104629517\n",
      "0.15700982511043549\n",
      "0.1560421884059906\n",
      "0.15508392453193665\n",
      "0.15413634479045868\n",
      "epoch: 141, accuracy: 0.9562000036239624\n",
      "0.15319764614105225\n",
      "0.15226760506629944\n",
      "0.1513466089963913\n",
      "0.1504340022802353\n",
      "0.14952947199344635\n",
      "0.1486329436302185\n",
      "0.14774490892887115\n",
      "0.14686577022075653\n",
      "0.14599448442459106\n",
      "0.14513073861598969\n",
      "epoch: 151, accuracy: 0.9584999680519104\n",
      "0.14427441358566284\n",
      "0.14342571794986725\n",
      "0.14258542656898499\n",
      "0.1417527198791504\n",
      "0.14092783629894257\n",
      "0.14011037349700928\n",
      "0.1393008828163147\n",
      "0.138499915599823\n",
      "0.13770736753940582\n",
      "0.13692249357700348\n",
      "epoch: 161, accuracy: 0.9599999785423279\n",
      "0.1361442655324936\n",
      "0.13537251949310303\n",
      "0.1346076875925064\n",
      "0.13384893536567688\n",
      "0.13309700787067413\n",
      "0.1323518007993698\n",
      "0.13161303102970123\n",
      "0.13088040053844452\n",
      "0.13015349209308624\n",
      "0.12943223118782043\n",
      "epoch: 171, accuracy: 0.9614999890327454\n",
      "0.12871772050857544\n",
      "0.12800903618335724\n",
      "0.1273062527179718\n",
      "0.12660929560661316\n",
      "0.12591780722141266\n",
      "0.12523192167282104\n",
      "0.12455156445503235\n",
      "0.12387610226869583\n",
      "0.12320593744516373\n",
      "0.12254077941179276\n",
      "epoch: 181, accuracy: 0.9631999731063843\n",
      "0.12188049405813217\n",
      "0.12122564017772675\n",
      "0.1205763891339302\n",
      "0.11993244290351868\n",
      "0.11929372698068619\n",
      "0.11865918338298798\n",
      "0.11802872270345688\n",
      "0.11740278452634811\n",
      "0.11678140610456467\n",
      "0.11616399139165878\n",
      "epoch: 191, accuracy: 0.9643999934196472\n",
      "0.1155511736869812\n",
      "0.11494361609220505\n",
      "0.11434081941843033\n",
      "0.11374329775571823\n",
      "0.11315032094717026\n",
      "0.11256197839975357\n",
      "0.11197807639837265\n",
      "0.11139848083257675\n",
      "0.11082252115011215\n",
      "0.11025029420852661\n",
      "epoch: 201, accuracy: 0.9649999737739563\n",
      "0.10968202352523804\n",
      "0.10911735892295837\n",
      "0.10855652391910553\n",
      "0.10799964517354965\n",
      "0.10744680464267731\n",
      "0.10689772665500641\n",
      "0.10635258257389069\n",
      "0.10581144690513611\n",
      "0.10527385026216507\n",
      "0.10473977774381638\n",
      "epoch: 211, accuracy: 0.9666999578475952\n",
      "0.10420923680067062\n",
      "0.10368240624666214\n",
      "0.10315912961959839\n",
      "0.10263917595148087\n",
      "0.10212287306785583\n",
      "0.10161018371582031\n",
      "0.10110075026750565\n",
      "0.10059473663568497\n",
      "0.10009163618087769\n",
      "0.09959188848733902\n",
      "epoch: 221, accuracy: 0.9678999781608582\n",
      "0.09909532964229584\n",
      "0.09860203415155411\n",
      "0.09811223298311234\n",
      "0.09762564301490784\n",
      "0.097142294049263\n",
      "0.09666202962398529\n",
      "0.09618517011404037\n",
      "0.09571115672588348\n",
      "0.09524013847112656\n",
      "0.09477216005325317\n",
      "epoch: 231, accuracy: 0.9685999751091003\n",
      "0.09430734068155289\n",
      "0.09384579211473465\n",
      "0.09338710457086563\n",
      "0.09293121844530106\n",
      "epoch: 235, accuracy: 0.9685999751091003\n",
      "0.09247834235429764\n",
      "0.09202823042869568\n",
      "0.09158076345920563\n",
      "0.09113606810569763\n",
      "0.09069420397281647\n",
      "0.09025539457798004\n",
      "epoch: 241, accuracy: 0.9688999652862549\n",
      "0.08981993794441223\n",
      "0.0893874317407608\n",
      "0.08895789086818695\n",
      "0.08853123337030411\n",
      "0.08810758590698242\n",
      "0.08768650144338608\n",
      "0.08726833015680313\n",
      "0.08685300499200821\n",
      "0.08644092082977295\n",
      "0.08603142946958542\n",
      "epoch: 251, accuracy: 0.9693999886512756\n",
      "0.08562500774860382\n",
      "0.08522088825702667\n",
      "0.08481946587562561\n",
      "0.08442065119743347\n",
      "0.08402447402477264\n",
      "0.0836305171251297\n",
      "0.08323891460895538\n",
      "0.08285009860992432\n",
      "0.08246345818042755\n",
      "0.08207940310239792\n",
      "epoch: 261, accuracy: 0.9705999493598938\n",
      "0.08169753104448318\n",
      "0.08131802082061768\n",
      "0.08094079047441483\n",
      "0.08056601881980896\n",
      "0.08019333332777023\n",
      "0.0798230767250061\n",
      "0.07945539057254791\n",
      "0.07908983528614044\n",
      "0.07872690260410309\n",
      "0.07836621254682541\n",
      "epoch: 271, accuracy: 0.9711999893188477\n",
      "0.07800795882940292\n",
      "0.0776515081524849\n",
      "0.07729768753051758\n",
      "0.07694578170776367\n",
      "0.07659585773944855\n",
      "0.07624779641628265\n",
      "0.07590219378471375\n",
      "0.07555906474590302\n",
      "0.07521802932024002\n",
      "0.07487893104553223\n",
      "epoch: 281, accuracy: 0.9717999696731567\n",
      "0.07454204559326172\n",
      "0.07420655339956284\n",
      "0.0738731399178505\n",
      "0.07354140281677246\n",
      "0.07321208715438843\n",
      "0.07288472354412079\n",
      "0.07255935668945312\n",
      "0.07223578542470932\n",
      "0.07191397994756699\n",
      "0.07159414142370224\n",
      "epoch: 291, accuracy: 0.972000002861023\n",
      "0.07127616554498672\n",
      "0.07096026092767715\n",
      "0.07064590603113174\n",
      "0.0703335627913475\n",
      "0.0700230598449707\n",
      "0.06971441954374313\n",
      "0.06940705329179764\n",
      "0.06910137832164764\n",
      "0.0687970221042633\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(300):\n",
    "    model.train()\n",
    "    for iter_idx, (data, label) in enumerate(train_loader):\n",
    "        data, label = data.to(device), label.to(device)\n",
    "        optim.zero_grad()\n",
    "        preds = model(data)\n",
    "        loss = criterion(preds, label)\n",
    "        print(loss.item())\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "\n",
    "        if epoch == 0 or (epoch % 10) == 0 or epoch == 234:\n",
    "            model.eval()\n",
    "            tp_count = 0\n",
    "            val_size = 0\n",
    "            with torch.no_grad():\n",
    "                for val_iter_idx, (val_data, val_label) in enumerate(test_loader):\n",
    "                    val_data, val_label = val_data.to(device), val_label.to(device)\n",
    "                    val_preds = model(val_data)\n",
    "                    val_preds = torch.argmax(val_preds, dim=1)\n",
    "                    tp_count += torch.count_nonzero(val_label == val_preds)\n",
    "                    val_size += val_data.shape[0]\n",
    "                print('epoch: {}, accuracy: {}'.format(epoch+1, (tp_count / val_size).item()))\n",
    "            model.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-stable",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}