Shakespeare/FedAvg.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "FedAvg.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "Y-ZegBArf4xQ",
        "cce_-qnxhD4n",
        "XelbyPsDlfgb",
        "IOBblyFGlwlU",
        "sWVOxcAao2_t",
        "vFFAfTOwpk4j",
        "c640e4NnpksE",
        "3crFDN0xqGu6",
        "YXtGLkoAqLIW"
      ],
      "toc_visible": true,
      "authorship_tag": "ABX9TyODmfkWrZZ1tNBNYaII6os2",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tushar-semwal/fedperf/blob/main/Santiago/Shakespeare/FedAvg.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZXnVpPCFfQkm"
      },
      "source": [
        "# FedPerf - Shakespeare + FedAvg algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TWDjzbCzfUFt"
      },
      "source": [
        "## Setup & Dependencies Installation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z_f-rWudW0Fo"
      },
      "source": [
        "%%capture\n",
        "!pip install torchsummaryX unidecode"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3TmMInb_fnw_"
      },
      "source": [
        "%load_ext tensorboard\n",
        "\n",
        "import copy\n",
        "from functools import reduce\n",
        "import json\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "import pandas as pd\n",
        "import pickle\n",
        "import random\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.preprocessing import StandardScaler, MinMaxScaler\n",
        "import time\n",
        "import torch\n",
        "from torch.autograd import Variable\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data.dataloader import DataLoader\n",
        "from torch.utils.data.sampler import Sampler\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from torchsummary import summary\n",
        "from torchsummaryX import summary as summaryx\n",
        "from torchvision import transforms, utils, datasets\n",
        "from tqdm.notebook import tqdm\n",
        "from unidecode import unidecode\n",
        "\n",
        "%matplotlib inline\n",
        "\n",
        "# Check assigned GPU\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)\n",
        "\n",
        "# set manual seed for reproducibility\n",
        "RANDOM_SEED = 42\n",
        "\n",
        "# general reproducibility\n",
        "random.seed(RANDOM_SEED)\n",
        "np.random.seed(RANDOM_SEED)\n",
        "torch.manual_seed(RANDOM_SEED)\n",
        "torch.cuda.manual_seed(RANDOM_SEED)\n",
        "\n",
        "# gpu training specific\n",
        "torch.backends.cudnn.deterministic = True\n",
        "torch.backends.cudnn.benchmark = False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Mfqv6uHfwh-"
      },
      "source": [
        "## Mount GDrive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "75qbJwxsGj-k"
      },
      "source": [
        "BASE_DIR = '/content/drive/MyDrive/FedPerf/shakespeare/FedAvg'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ost-6lXSfveC"
      },
      "source": [
        "try:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    os.makedirs(BASE_DIR, exist_ok=True)\n",
        "except:\n",
        "    print(\"WARNING: Results won't be stored on GDrive\")\n",
        "    BASE_DIR = './'\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y-ZegBArf4xQ"
      },
      "source": [
        "## Loading Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hf03LRxof7Zj"
      },
      "source": [
        "!rm -Rf data\n",
        "!mkdir -p data scripts"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ngygA4-Fgobx"
      },
      "source": [
        "GENERATE_DATASET = False  # If False, download the dataset provided by the q-FFL paper\n",
        "DATA_DIR = 'data/'\n",
        "# Dataset generation params\n",
        "SAMPLES_FRACTION = 1.  # If using an already generated dataset\n",
        "# SAMPLES_FRACTION = 0.2  # Fraction of total samples in the dataset - FedProx default script\n",
        "# SAMPLES_FRACTION = 0.05  # Fraction of total samples in the dataset - qFFL\n",
        "TRAIN_FRACTION = 0.8  # Train set size\n",
        "MIN_SAMPLES = 0  # Min samples per client (for filtering purposes) - FedProx\n",
        "# MIN_SAMPLES = 64  # Min samples per client (for filtering purposes) - qFFL"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nUmwJgJygoYD"
      },
      "source": [
        "# Download raw dataset\n",
        "# !wget https://ocw.mit.edu/ans7870/6/6.006/s08/lecturenotes/files/t8.shakespeare.txt -O data/shakespeare.txt\n",
        "!wget --adjust-extension http://www.gutenberg.org/files/100/100-0.txt -O data/shakespeare.txt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4dCvx80BgoVr"
      },
      "source": [
        "if not GENERATE_DATASET:\n",
        "    !rm -Rf data/train data/test\n",
        "    !gdown --id 1n46Mftp3_ahRi1Z6jYhEriyLtdRDS1tD  # Download Shakespeare dataset used by the FedProx paper\n",
        "    !unzip shakespeare.zip\n",
        "    !mv -f shakespeare_paper/train data/\n",
        "    !mv -f shakespeare_paper/test data/\n",
        "    !rm -R shakespeare_paper/ shakespeare.zip\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a4pzFvPvhQhq"
      },
      "source": [
        "corpus = []\n",
        "with open('data/shakespeare.txt', 'r') as f:\n",
        "    data = list(unidecode(f.read()))\n",
        "    corpus = list(set(list(data)))\n",
        "print('Corpus Length:', len(corpus))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cce_-qnxhD4n"
      },
      "source": [
        "#### Dataset Preprocessing script"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rt13M4IcgoTV"
      },
      "source": [
        "%%capture\n",
        "if GENERATE_DATASET:\n",
        "    # Download dataset generation scripts\n",
        "    !wget https://raw.githubusercontent.com/ml-lab/FedProx/master/data/shakespeare/preprocess/preprocess_shakespeare.py -O scripts/preprocess_shakespeare.py\n",
        "    !wget https://raw.githubusercontent.com/ml-lab/FedProx/master/data/shakespeare/preprocess/shake_utils.py -O scripts/shake_utils.py\n",
        "    !wget https://raw.githubusercontent.com/ml-lab/FedProx/master/data/shakespeare/preprocess/gen_all_data.py -O scripts/gen_all_data.py\n",
        "\n",
        "    # Download data preprocessing scripts\n",
        "    !wget https://raw.githubusercontent.com/ml-lab/FedProx/master/utils/sample.py -O scripts/sample.py\n",
        "    !wget https://raw.githubusercontent.com/ml-lab/FedProx/master/utils/remove_users.py -O scripts/remove_users.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EIEyRW27goPo"
      },
      "source": [
        "# Running scripts\n",
        "if GENERATE_DATASET:\n",
        "    !mkdir -p data/raw_data data/all_data data/train data/test\n",
        "    !python scripts/preprocess_shakespeare.py data/shakespeare.txt data/raw_data\n",
        "    !python scripts/gen_all_data.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mq8V6v_4hhhD"
      },
      "source": [
        "#### Dataset class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H2SjEBKoWDxv"
      },
      "source": [
        "class ShakespeareDataset(Dataset):\n",
        "    def __init__(self, x, y, corpus, seq_length):\n",
        "        self.x = x\n",
        "        self.y = y\n",
        "        self.corpus = corpus\n",
        "        self.corpus_size = len(self.corpus)\n",
        "        super(ShakespeareDataset, self).__init__()\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.x)\n",
        "\n",
        "    def __repr__(self):\n",
        "        return f'{self.__class__} - (length: {self.__len__()})'\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        input_seq = self.x[i]\n",
        "        next_char = self.y[i]\n",
        "        # print('\\tgetitem', i, input_seq, next_char)\n",
        "        input_value = self.text2charindxs(input_seq)\n",
        "        target_value = self.get_label_from_char(next_char)\n",
        "        return input_value, target_value\n",
        "\n",
        "    def text2charindxs(self, text):\n",
        "        tensor = torch.zeros(len(text), dtype=torch.int32)\n",
        "        for i, c in enumerate(text):\n",
        "            tensor[i] = self.get_label_from_char(c)\n",
        "        return tensor\n",
        "\n",
        "    def get_label_from_char(self, c):\n",
        "        return self.corpus.index(c)\n",
        "\n",
        "    def get_char_from_label(self, l):\n",
        "        return self.corpus[l]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9fgJtS62lYAN"
      },
      "source": [
        "##### Federated Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5DqL5pTmgn5X"
      },
      "source": [
        "class ShakespeareFedDataset(ShakespeareDataset):\n",
        "    def __init__(self, x, y, corpus, seq_length):\n",
        "        super(ShakespeareFedDataset, self).__init__(x, y, corpus, seq_length)\n",
        "\n",
        "    def dataloader(self, batch_size, shuffle=True):\n",
        "        return DataLoader(self,\n",
        "                          batch_size=batch_size,\n",
        "                          shuffle=shuffle,\n",
        "                          num_workers=0)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XelbyPsDlfgb"
      },
      "source": [
        "## Partitioning & Data Loaders"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IOBblyFGlwlU"
      },
      "source": [
        "### IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cSZFWKmsgn1p"
      },
      "source": [
        "def iid_partition_(dataset, clients):\n",
        "  \"\"\"\n",
        "  I.I.D paritioning of data over clients\n",
        "  Shuffle the data\n",
        "  Split it between clients\n",
        "  \n",
        "  params:\n",
        "    - dataset (torch.utils.Dataset): Dataset\n",
        "    - clients (int): Number of Clients to split the data between\n",
        "\n",
        "  returns:\n",
        "    - Dictionary of image indexes for each client\n",
        "  \"\"\"\n",
        "\n",
        "  num_items_per_client = int(len(dataset)/clients)\n",
        "  client_dict = {}\n",
        "  image_idxs = [i for i in range(len(dataset))]\n",
        "\n",
        "  for i in range(clients):\n",
        "    client_dict[i] = set(np.random.choice(image_idxs, num_items_per_client, replace=False))\n",
        "    image_idxs = list(set(image_idxs) - client_dict[i])\n",
        "\n",
        "  return client_dict"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-lGwDyhSll9h"
      },
      "source": [
        "def iid_partition(corpus, seq_length=80, val_split=False):\n",
        "\n",
        "    train_file = [os.path.join(DATA_DIR, 'train', f) for f in os.listdir(f'{DATA_DIR}/train') if f.endswith('.json')][0]\n",
        "    test_file = [os.path.join(DATA_DIR, 'test', f) for f in os.listdir(f'{DATA_DIR}/test') if f.endswith('.json')][0]\n",
        "\n",
        "    with open(train_file, 'r') as file:\n",
        "        data_train = json.loads(unidecode(file.read()))\n",
        "\n",
        "    with open(test_file, 'r') as file:\n",
        "        data_test = json.loads(unidecode(file.read()))\n",
        "\n",
        "    \n",
        "    total_samples_train = sum(data_train['num_samples'])\n",
        "\n",
        "    data_dict = {}\n",
        "\n",
        "    x_train, y_train = [], []\n",
        "    x_test, y_test = [], []\n",
        "    # x_val, y_val = [], []\n",
        "\n",
        "    users = list(zip(data_train['users'], data_train['num_samples']))\n",
        "    # random.shuffle(users)\n",
        "\n",
        "\n",
        "\n",
        "    total_samples = int(sum(data_train['num_samples']) * SAMPLES_FRACTION)\n",
        "    print('Objective', total_samples, '/', sum(data_train['num_samples']))\n",
        "    sample_count = 0\n",
        "    \n",
        "    for i, (author_id, samples) in enumerate(users):\n",
        "\n",
        "        if sample_count >= total_samples:\n",
        "            print('Max samples reached', sample_count, '/', total_samples)\n",
        "            break\n",
        "\n",
        "        if samples < MIN_SAMPLES: # or data_train['num_samples'][i] > 10000:\n",
        "            print('SKIP', author_id, samples)\n",
        "            continue\n",
        "        else:\n",
        "            udata_train = data_train['user_data'][author_id]\n",
        "            max_samples = samples if (sample_count + samples) <= total_samples else (sample_count + samples - total_samples) \n",
        "            \n",
        "            sample_count += max_samples\n",
        "            # print('sample_count', sample_count)\n",
        "\n",
        "            x_train.extend(data_train['user_data'][author_id]['x'][:max_samples])\n",
        "            y_train.extend(data_train['user_data'][author_id]['y'][:max_samples])\n",
        "\n",
        "            author_data = data_test['user_data'][author_id]\n",
        "            test_size = int(len(author_data['x']) * SAMPLES_FRACTION)\n",
        "\n",
        "            if val_split:\n",
        "                x_test.extend(author_data['x'][:int(test_size / 2)])\n",
        "                y_test.extend(author_data['y'][:int(test_size / 2)])\n",
        "                # x_val.extend(author_data['x'][int(test_size / 2):])\n",
        "                # y_val.extend(author_data['y'][int(test_size / 2):int(test_size)])\n",
        "\n",
        "            else:\n",
        "                x_test.extend(author_data['x'][:int(test_size)])\n",
        "                y_test.extend(author_data['y'][:int(test_size)])\n",
        "\n",
        "    train_ds = ShakespeareDataset(x_train, y_train, corpus, seq_length)\n",
        "    test_ds = ShakespeareDataset(x_test, y_test, corpus, seq_length)\n",
        "    # val_ds = ShakespeareDataset(x_val, y_val, corpus, seq_length)\n",
        "\n",
        "    data_dict = iid_partition_(train_ds, clients=len(users))\n",
        "\n",
        "    return train_ds, data_dict, test_ds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFvc8mLoouKa"
      },
      "source": [
        "### Non-IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GZ76WsCZot9s"
      },
      "source": [
        "def noniid_partition(corpus, seq_length=80, val_split=False):\n",
        "\n",
        "    train_file = [os.path.join(DATA_DIR, 'train', f) for f in os.listdir(f'{DATA_DIR}/train') if f.endswith('.json')][0]\n",
        "    test_file = [os.path.join(DATA_DIR, 'test', f) for f in os.listdir(f'{DATA_DIR}/test') if f.endswith('.json')][0]\n",
        "\n",
        "    with open(train_file, 'r') as file:\n",
        "        data_train = json.loads(unidecode(file.read()))\n",
        "\n",
        "    with open(test_file, 'r') as file:\n",
        "        data_test = json.loads(unidecode(file.read()))\n",
        "\n",
        "    \n",
        "    total_samples_train = sum(data_train['num_samples'])\n",
        "\n",
        "    data_dict = {}\n",
        "\n",
        "    x_test, y_test = [], []\n",
        "\n",
        "    users = list(zip(data_train['users'], data_train['num_samples']))\n",
        "    # random.shuffle(users)\n",
        "\n",
        "    total_samples = int(sum(data_train['num_samples']) * SAMPLES_FRACTION)\n",
        "    print('Objective', total_samples, '/', sum(data_train['num_samples']))\n",
        "    sample_count = 0\n",
        "    \n",
        "    for i, (author_id, samples) in enumerate(users):\n",
        "\n",
        "        if sample_count >= total_samples:\n",
        "            print('Max samples reached', sample_count, '/', total_samples)\n",
        "            break\n",
        "\n",
        "        if samples < MIN_SAMPLES: # or data_train['num_samples'][i] > 10000:\n",
        "            print('SKIP', author_id, samples)\n",
        "            continue\n",
        "        else:\n",
        "            udata_train = data_train['user_data'][author_id]\n",
        "            max_samples = samples if (sample_count + samples) <= total_samples else (sample_count + samples - total_samples) \n",
        "            \n",
        "            sample_count += max_samples\n",
        "            # print('sample_count', sample_count)\n",
        "\n",
        "            x_train = data_train['user_data'][author_id]['x'][:max_samples]\n",
        "            y_train = data_train['user_data'][author_id]['y'][:max_samples]\n",
        "\n",
        "            train_ds = ShakespeareFedDataset(x_train, y_train, corpus, seq_length)\n",
        "\n",
        "            x_val, y_val = None, None\n",
        "            val_ds = None\n",
        "            author_data = data_test['user_data'][author_id]\n",
        "            test_size = int(len(author_data['x']) * SAMPLES_FRACTION)\n",
        "            if val_split:\n",
        "                x_test += author_data['x'][:int(test_size / 2)]\n",
        "                y_test += author_data['y'][:int(test_size / 2)]\n",
        "                x_val = author_data['x'][int(test_size / 2):]\n",
        "                y_val = author_data['y'][int(test_size / 2):int(test_size)]\n",
        "\n",
        "                val_ds = ShakespeareFedDataset(x_val, y_val, corpus, seq_length)\n",
        "\n",
        "            else:\n",
        "                x_test += author_data['x'][:int(test_size)]\n",
        "                y_test += author_data['y'][:int(test_size)]\n",
        "\n",
        "            data_dict[author_id] = {\n",
        "                'train_ds': train_ds,\n",
        "                'val_ds': val_ds\n",
        "            }\n",
        "\n",
        "    test_ds = ShakespeareFedDataset(x_test, y_test, corpus, seq_length)\n",
        "\n",
        "    return data_dict, test_ds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sWVOxcAao2_t"
      },
      "source": [
        "## Models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gQQQ2mLeo6EA"
      },
      "source": [
        "### Shakespeare LSTM"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2mGXTrXRot7R"
      },
      "source": [
        "class ShakespeareLSTM(nn.Module):\n",
        "    \"\"\"\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, input_dim, embedding_dim, hidden_dim, classes, lstm_layers=2, dropout=0.1, batch_first=True):\n",
        "        super(ShakespeareLSTM, self).__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.embedding_dim = embedding_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.classes = classes\n",
        "        self.no_layers = lstm_layers\n",
        "        \n",
        "        self.embedding = nn.Embedding(num_embeddings=self.classes,\n",
        "                                      embedding_dim=self.embedding_dim)\n",
        "        self.lstm = nn.LSTM(input_size=self.embedding_dim, \n",
        "                            hidden_size=self.hidden_dim,\n",
        "                            num_layers=self.no_layers,\n",
        "                            batch_first=batch_first, \n",
        "                            dropout=dropout if self.no_layers > 1 else 0.)\n",
        "        self.fc = nn.Linear(hidden_dim, self.classes)\n",
        "\n",
        "    def forward(self, x, hc=None):\n",
        "        batch_size = x.size(0)\n",
        "        x_emb = self.embedding(x)\n",
        "        out, (ht, ct) = self.lstm(x_emb.view(batch_size, -1, self.embedding_dim), hc)\n",
        "        dense = self.fc(ht[-1])\n",
        "        return dense\n",
        "    \n",
        "    def init_hidden(self, batch_size):\n",
        "        return (Variable(torch.zeros(self.no_layers, batch_size, self.hidden_dim)),\n",
        "                Variable(torch.zeros(self.no_layers, batch_size, self.hidden_dim)))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5QsuJlVipMc8"
      },
      "source": [
        "#### Model Summary"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n_Vb0BYpot5I"
      },
      "source": [
        "batch_size = 10\n",
        "seq_length = 80 # mcmahan17a, fedprox, qFFL\n",
        "\n",
        "shakespeare_lstm = ShakespeareLSTM(input_dim=seq_length,  \n",
        "                                   embedding_dim=8,  # mcmahan17a, fedprox, qFFL\n",
        "                                   hidden_dim=256,  # mcmahan17a, fedprox impl\n",
        "                                #    hidden_dim=100,  # fedprox paper\n",
        "                                   classes=len(corpus),\n",
        "                                   lstm_layers=2,\n",
        "                                   dropout=0.1,  # TODO:\n",
        "                                   batch_first=True\n",
        "                                   )\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "  shakespeare_lstm.cuda()\n",
        "\n",
        "\n",
        "\n",
        "hc = shakespeare_lstm.init_hidden(batch_size)\n",
        "\n",
        "x_sample = torch.zeros((batch_size, seq_length),\n",
        "                       dtype=torch.long,\n",
        "                       device=(torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')))\n",
        "\n",
        "x_sample[0][0] = 1\n",
        "x_sample\n",
        "\n",
        "print(\"\\nShakespeare LSTM SUMMARY\")\n",
        "print(summaryx(shakespeare_lstm, x_sample))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qn7egnzTpeks"
      },
      "source": [
        "## FedAvg Algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vFFAfTOwpk4j"
      },
      "source": [
        "### Plot Utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oyYjWa6IpnTY"
      },
      "source": [
        "from sklearn.metrics import f1_score"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "367THsiTpo-C"
      },
      "source": [
        "def plot_scores(history, exp_id, title, suffix):\n",
        "    accuracies = [x['accuracy'] for x in history]\n",
        "    f1_macro = [x['f1_macro'] for x in history]\n",
        "    f1_weighted = [x['f1_weighted'] for x in history]\n",
        "\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.plot(accuracies, 'tab:orange')\n",
        "    ax.set(xlabel='Rounds', ylabel='Test Accuracy', title=title)\n",
        "    ax.grid()\n",
        "    fig.savefig(f'{BASE_DIR}/{exp_id}/Test_Accuracy_{suffix}.jpg', format='jpg', dpi=300)\n",
        "    plt.show()\n",
        "\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.plot(f1_macro, 'tab:orange')\n",
        "    ax.set(xlabel='Rounds', ylabel='Test F1 (macro)', title=title)\n",
        "    ax.grid()\n",
        "    fig.savefig(f'{BASE_DIR}/{exp_id}/Test_F1_Macro_{suffix}.jpg', format='jpg')\n",
        "    plt.show()\n",
        "\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.plot(f1_weighted, 'tab:orange')\n",
        "    ax.set(xlabel='Rounds', ylabel='Test F1 (weighted)', title=title)\n",
        "    ax.grid()\n",
        "    fig.savefig(f'{BASE_DIR}/{exp_id}/Test_F1_Weighted_{suffix}.jpg', format='jpg')\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def plot_losses(history, exp_id, title, suffix):\n",
        "    val_losses = [x['loss'] for x in history]\n",
        "    train_losses = [x['train_loss'] for x in history]\n",
        "\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.plot(train_losses, 'tab:orange')\n",
        "    ax.set(xlabel='Rounds', ylabel='Train Loss', title=title)\n",
        "    ax.grid()\n",
        "    fig.savefig(f'{BASE_DIR}/{exp_id}/Train_Loss_{suffix}.jpg', format='jpg')\n",
        "    plt.show()\n",
        "\n",
        "    fig, ax = plt.subplots()\n",
        "    ax.plot(val_losses, 'tab:orange')\n",
        "    ax.set(xlabel='Rounds', ylabel='Test Loss', title=title)\n",
        "    ax.grid()\n",
        "    fig.savefig(f'{BASE_DIR}/{exp_id}/Test_Loss_{suffix}.jpg', format='jpg')\n",
        "    plt.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c640e4NnpksE"
      },
      "source": [
        "### Systems Heterogeneity Simulations\n",
        "\n",
        "Generate epochs for selected clients based on percentage of devices that corresponds to heterogeneity. \n",
        "\n",
        "Assign x number of epochs (chosen unifirmly at random between [1, E]) to 0%, 50% or 90% of the selected devices, respectively. Settings where 0% devices perform fewer than E epochs of work correspond to the environments without system heterogeneity, while 90% of the devices sending their partial solutions corresponds to highly heterogenous system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zuEZYnl5ot2m"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VQ9PZM0Gp9ve"
      },
      "source": [
        "### Local Training (Client Update)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EDJFltwdotzZ"
      },
      "source": [
        "class CustomDataset(Dataset):\n",
        "  def __init__(self, dataset, idxs):\n",
        "      self.dataset = dataset\n",
        "      self.idxs = list(idxs)\n",
        "\n",
        "  def __len__(self):\n",
        "      return len(self.idxs)\n",
        "\n",
        "  def __getitem__(self, item):\n",
        "      data, label = self.dataset[self.idxs[item]]\n",
        "      return data, label"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HtRzU5Yepddq"
      },
      "source": [
        "class ClientUpdate(object):\n",
        "  def __init__(self, dataset, batchSize, learning_rate, epochs, idxs, mu, algorithm):\n",
        "    # self.train_loader = DataLoader(CustomDataset(dataset, idxs), batch_size=batchSize, shuffle=True)\n",
        "    if hasattr(dataset, 'dataloader'):\n",
        "        self.train_loader = dataset.dataloader(batch_size=batch_size, shuffle=True)\n",
        "    else:\n",
        "        self.train_loader = DataLoader(CustomDataset(dataset, idxs), batch_size=batch_size, shuffle=True)\n",
        "\n",
        "    self.algorithm = algorithm\n",
        "    self.learning_rate = learning_rate\n",
        "    self.epochs = epochs\n",
        "    self.mu = mu\n",
        "\n",
        "  def train(self, model):\n",
        "    # print(\"Client training for {} epochs.\".format(self.epochs))\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    proximal_criterion = nn.MSELoss(reduction='mean')\n",
        "    optimizer = torch.optim.SGD(model.parameters(), lr=self.learning_rate, momentum=0.5)\n",
        "\n",
        "    # use the weights of global model for proximal term calculation\n",
        "    global_model = copy.deepcopy(model)\n",
        "\n",
        "    # calculate local training time\n",
        "    start_time = time.time()\n",
        "\n",
        "\n",
        "    e_loss = []\n",
        "    for epoch in range(1, self.epochs+1):\n",
        "\n",
        "      train_loss = 0.0\n",
        "\n",
        "      model.train()\n",
        "      for data, labels in self.train_loader:\n",
        "\n",
        "        if torch.cuda.is_available():\n",
        "          data, labels = data.cuda(), labels.cuda()\n",
        "\n",
        "        # clear the gradients\n",
        "        optimizer.zero_grad()\n",
        "        # make a forward pass\n",
        "        output = model(data)\n",
        "\n",
        "        # calculate the loss + the proximal term\n",
        "        _, pred = torch.max(output, 1)\n",
        "\n",
        "        if self.algorithm == 'fedprox':\n",
        "          proximal_term = 0.0\n",
        "\n",
        "          # iterate through the current and global model parameters\n",
        "          for w, w_t in zip(model.parameters(), global_model.parameters()) :\n",
        "            # update the proximal term \n",
        "            #proximal_term += torch.sum(torch.abs((w-w_t)**2))\n",
        "            proximal_term += (w-w_t).norm(2)\n",
        "\n",
        "          loss = criterion(output, labels) + (self.mu/2)*proximal_term\n",
        "        else:\n",
        "          loss = criterion(output, labels)\n",
        "    \n",
        "        # do a backwards pass\n",
        "        loss.backward()\n",
        "        # perform a single optimization step\n",
        "        optimizer.step()\n",
        "        # update training loss\n",
        "        train_loss += loss.item()*data.size(0)\n",
        "\n",
        "      # average losses\n",
        "      train_loss = train_loss/len(self.train_loader.dataset)\n",
        "      e_loss.append(train_loss)\n",
        "\n",
        "    total_loss = sum(e_loss)/len(e_loss)\n",
        "\n",
        "    return model.state_dict(), total_loss, (time.time() - start_time)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3crFDN0xqGu6"
      },
      "source": [
        "### Server Side Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c085xSOoqEHk"
      },
      "source": [
        "def training(model, rounds, batch_size, lr, ds, data_dict, test_ds, C, K, E, mu, percentage, plt_title, plt_color, target_test_accuracy,\n",
        "             classes, algorithm=\"fedprox\", history=[], eval_every=1, tb_logger=None):\n",
        "  \"\"\"\n",
        "  Function implements the Federated Averaging Algorithm from the FedAvg paper.\n",
        "  Specifically, this function is used for the server side training and weight update\n",
        "\n",
        "  Params:\n",
        "    - model:           PyTorch model to train\n",
        "    - rounds:          Number of communication rounds for the client update\n",
        "    - batch_size:      Batch size for client update training\n",
        "    - lr:              Learning rate used for client update training\n",
        "    - ds:              Dataset used for training\n",
        "    - data_dict:       Type of data partition used for training (IID or non-IID)\n",
        "    - test_data_dict:  Data used for testing the model\n",
        "    - C:               Fraction of clients randomly chosen to perform computation on each round\n",
        "    - K:               Total number of clients\n",
        "    - E:               Number of training passes each client makes over its local dataset per round\n",
        "    - mu:              proximal term constant\n",
        "    - percentage:      percentage of selected client to have fewer than E epochs\n",
        "  Returns:\n",
        "    - model:           Trained model on the server\n",
        "  \"\"\"\n",
        "\n",
        "  start = time.time()\n",
        "\n",
        "  # global model weights\n",
        "  global_weights = model.state_dict()\n",
        "\n",
        "  # training loss\n",
        "  train_loss = []\n",
        "\n",
        "  # test accuracy\n",
        "  test_acc = []\n",
        "\n",
        "  # store last loss for convergence\n",
        "  last_loss = 0.0\n",
        "\n",
        "  # total time taken \n",
        "  total_time = 0\n",
        "\n",
        "  print(f\"System heterogeneity set to {percentage}% stragglers.\\n\")\n",
        "  print(f\"Picking {max(int(C*K),1 )} random clients per round.\\n\")\n",
        "\n",
        "  users_id = list(data_dict.keys())\n",
        "\n",
        "  for curr_round in range(1, rounds+1):\n",
        "    w, local_loss, lst_local_train_time = [], [], []\n",
        "\n",
        "    m = max(int(C*K), 1)\n",
        "\n",
        "    heterogenous_epoch_list = GenerateLocalEpochs(percentage, size=m, max_epochs=E)\n",
        "    heterogenous_epoch_list = np.array(heterogenous_epoch_list)\n",
        "    # print('heterogenous_epoch_list', len(heterogenous_epoch_list))\n",
        "\n",
        "    S_t = np.random.choice(range(K), m, replace=False)\n",
        "    S_t = np.array(S_t)\n",
        "    print('Clients: {}/{} -> {}'.format(len(S_t), K, S_t))\n",
        "    \n",
        "    # For Federated Averaging, drop all the clients that are stragglers\n",
        "    if algorithm == 'fedavg':\n",
        "      stragglers_indices = np.argwhere(heterogenous_epoch_list < E)\n",
        "      heterogenous_epoch_list = np.delete(heterogenous_epoch_list, stragglers_indices)\n",
        "      S_t = np.delete(S_t, stragglers_indices)\n",
        "\n",
        "    # for _, (k, epoch) in tqdm(enumerate(zip(S_t, heterogenous_epoch_list))):\n",
        "    for i in tqdm(range(len(S_t))):\n",
        "    #   print('k', k)\n",
        "      k = S_t[i]\n",
        "      epoch = heterogenous_epoch_list[i]\n",
        "      key = users_id[k]\n",
        "      ds_ = ds if ds else data_dict[key]['train_ds']\n",
        "      idxs = data_dict[key] if ds else None\n",
        "    #   print(f'Client {k}: {len(idxs) if idxs else len(ds_)} samples')\n",
        "      local_update = ClientUpdate(dataset=ds_, batchSize=batch_size, learning_rate=lr, epochs=epoch, idxs=idxs, mu=mu, algorithm=algorithm)\n",
        "      weights, loss, local_train_time = local_update.train(model=copy.deepcopy(model))\n",
        "    #   print(f'Local train time for {k} on {len(idxs) if idxs else len(ds_)} samples: {local_train_time}')\n",
        "    #   print(f'Local train time: {local_train_time}')\n",
        "\n",
        "      w.append(copy.deepcopy(weights))\n",
        "      local_loss.append(copy.deepcopy(loss))\n",
        "      lst_local_train_time.append(local_train_time)\n",
        "\n",
        "    # calculate time to update the global weights\n",
        "    global_start_time = time.time()\n",
        "\n",
        "    # updating the global weights\n",
        "    weights_avg = copy.deepcopy(w[0])\n",
        "    for k in weights_avg.keys():\n",
        "      for i in range(1, len(w)):\n",
        "        weights_avg[k] += w[i][k]\n",
        "\n",
        "      weights_avg[k] = torch.div(weights_avg[k], len(w))\n",
        "\n",
        "    global_weights = weights_avg\n",
        "\n",
        "    global_end_time = time.time()\n",
        "\n",
        "    # calculate total time \n",
        "    total_time += (global_end_time - global_start_time) + sum(lst_local_train_time)/len(lst_local_train_time)\n",
        "\n",
        "    # move the updated weights to our model state dict\n",
        "    model.load_state_dict(global_weights)\n",
        "\n",
        "    # loss\n",
        "    loss_avg = sum(local_loss) / len(local_loss)\n",
        "    print('Round: {}... \\tAverage Loss: {}'.format(curr_round, round(loss_avg, 3)))\n",
        "    train_loss.append(loss_avg)\n",
        "    if tb_logger:\n",
        "        tb_logger.add_scalar(f'Train/Loss', loss_avg, curr_round)\n",
        "\n",
        "    # testing\n",
        "    # if curr_round % eval_every == 0:\n",
        "    test_scores = testing(model, test_ds, batch_size * 2, nn.CrossEntropyLoss(), len(classes), classes)\n",
        "    test_scores['train_loss'] = loss_avg\n",
        "    test_loss, test_accuracy = test_scores['loss'], test_scores['accuracy']\n",
        "    history.append(test_scores)\n",
        "    \n",
        "    # print('Round: {}... \\tAverage Loss: {} \\tTest Loss: {} \\tTest Acc: {}'.format(curr_round, round(loss_avg, 3), round(test_loss, 3), round(test_accuracy, 3)))\n",
        "\n",
        "    if tb_logger:\n",
        "        tb_logger.add_scalar(f'Test/Loss', test_scores['loss'], curr_round)\n",
        "        tb_logger.add_scalars(f'Test/Scores', {\n",
        "            'accuracy': test_scores['accuracy'], 'f1_macro': test_scores['f1_macro'], 'f1_weighted': test_scores['f1_weighted']\n",
        "        }, curr_round)\n",
        "\n",
        "    test_acc.append(test_accuracy)\n",
        "    # break if we achieve the target test accuracy\n",
        "    if test_accuracy >= target_test_accuracy:\n",
        "      rounds = curr_round\n",
        "      break\n",
        "\n",
        "    # break if we achieve convergence, i.e., loss between two consecutive rounds is <0.0001\n",
        "    if algorithm == 'fedprox' and abs(loss_avg - last_loss) < 1e-5:\n",
        "      rounds = curr_round\n",
        "      break\n",
        "    \n",
        "    # update the last loss\n",
        "    last_loss = loss_avg\n",
        "\n",
        "  end = time.time()\n",
        "  \n",
        "  # plot train loss\n",
        "  fig, ax = plt.subplots()\n",
        "  x_axis = np.arange(1, rounds+1)\n",
        "  y_axis = np.array(train_loss)\n",
        "  ax.plot(x_axis, y_axis)\n",
        "\n",
        "  ax.set(xlabel='Number of Rounds', ylabel='Train Loss', title=plt_title)\n",
        "  ax.grid()\n",
        "  # fig.savefig(plt_title+'.jpg', format='jpg')\n",
        "\n",
        "  # plot test accuracy\n",
        "  fig1, ax1 = plt.subplots()\n",
        "  x_axis1 = np.arange(1, rounds+1)\n",
        "  y_axis1 = np.array(test_acc)\n",
        "  ax1.plot(x_axis1, y_axis1)\n",
        "\n",
        "  ax1.set(xlabel='Number of Rounds', ylabel='Test Accuracy', title=plt_title)\n",
        "  ax1.grid()\n",
        "  # fig1.savefig(plt_title+'-test.jpg', format='jpg')\n",
        "  \n",
        "  print(\"Training Done! Total time taken to Train: {}\".format(end-start))\n",
        "\n",
        "  return model, history"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YXtGLkoAqLIW"
      },
      "source": [
        "### Testing Loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dQJIJno4qKvc"
      },
      "source": [
        "def testing(model, dataset, bs, criterion, num_classes, classes, print_all=False):\n",
        "  #test loss \n",
        "  test_loss = 0.0\n",
        "  correct_class = list(0. for i in range(num_classes))\n",
        "  total_class = list(0. for i in range(num_classes))\n",
        "\n",
        "  test_loader = DataLoader(dataset, batch_size=bs)\n",
        "  l = len(test_loader)\n",
        "  model.eval()\n",
        "  print('running validation...')\n",
        "  for i, (data, labels) in enumerate(tqdm(test_loader)):\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "      data, labels = data.cuda(), labels.cuda()\n",
        "\n",
        "    output = model(data)\n",
        "    loss = criterion(output, labels)\n",
        "    test_loss += loss.item()*data.size(0)\n",
        "\n",
        "    _, pred = torch.max(output, 1)\n",
        "\n",
        "    # For F1Score\n",
        "    y_true = np.append(y_true, labels.data.view_as(pred).cpu().numpy()) if i != 0 else labels.data.view_as(pred).cpu().numpy()\n",
        "    y_hat = np.append(y_hat, pred.cpu().numpy()) if i != 0 else pred.cpu().numpy()\n",
        "\n",
        "    correct_tensor = pred.eq(labels.data.view_as(pred))\n",
        "    correct = np.squeeze(correct_tensor.numpy()) if not torch.cuda.is_available() else np.squeeze(correct_tensor.cpu().numpy())\n",
        "\n",
        "    #test accuracy for each object class\n",
        "    # for i in range(num_classes):\n",
        "    #   label = labels.data[i]\n",
        "    #   correct_class[label] += correct[i].item()\n",
        "    #   total_class[label] += 1\n",
        "\n",
        "    for i, lbl in enumerate(labels.data):\n",
        "    #   print('lbl', i, lbl)\n",
        "      correct_class[lbl] += correct.data[i]\n",
        "      total_class[lbl] += 1\n",
        "    \n",
        "  # avg test loss\n",
        "  test_loss = test_loss/len(test_loader.dataset)\n",
        "  print(\"Test Loss: {:.6f}\\n\".format(test_loss))\n",
        "\n",
        "  # Avg F1 Score\n",
        "  f1_macro = f1_score(y_true, y_hat, average='macro')\n",
        "  # F1-Score -> weigthed to consider class imbalance\n",
        "  f1_weighted =  f1_score(y_true, y_hat, average='weighted')\n",
        "  print(\"F1 Score: {:.6f} (macro) {:.6f} (weighted) %\\n\".format(f1_macro, f1_weighted))\n",
        "\n",
        "  # print test accuracy\n",
        "  if print_all:\n",
        "    for i in range(num_classes):\n",
        "        if total_class[i]>0:\n",
        "            print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % \n",
        "                    (classes[i], 100 * correct_class[i] / total_class[i],\n",
        "                    np.sum(correct_class[i]), np.sum(total_class[i])))\n",
        "        else:\n",
        "            print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))\n",
        "\n",
        "  overall_accuracy = np.sum(correct_class) / np.sum(total_class)\n",
        "\n",
        "  print('\\nFinal Test  Accuracy: {:.3f} ({}/{})'.format(overall_accuracy, np.sum(correct_class), np.sum(total_class)))\n",
        "\n",
        "  return {'loss': test_loss, 'accuracy': overall_accuracy, 'f1_macro': f1_macro, 'f1_weighted': f1_weighted}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uxqXLBd8qbC2"
      },
      "source": [
        "## Experiments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VRKlrkVHO8Na"
      },
      "source": [
        "FAIL-ON-PURPOSE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E2CfSkNVqKtL"
      },
      "source": [
        "seq_length = 80  # mcmahan17a, fedprox, qFFL\n",
        "embedding_dim = 8  # mcmahan17a, fedprox, qFFL\n",
        "# hidden_dim = 100  # fedprox paper\n",
        "hidden_dim = 256  # mcmahan17a, fedprox impl\n",
        "num_classes = len(corpus)\n",
        "classes = list(range(num_classes))\n",
        "lstm_layers = 2  # mcmahan17a, fedprox, qFFL\n",
        "dropout = 0.1  # TODO\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OvCgT_qbmFAO"
      },
      "source": [
        "class Hyperparameters():\n",
        "\n",
        "    def __init__(self, total_clients):\n",
        "        # number of training rounds\n",
        "        self.rounds = 50\n",
        "        # client fraction\n",
        "        self.C = 0.5\n",
        "        # number of clients\n",
        "        self.K = total_clients\n",
        "        # number of training passes on local dataset for each roung\n",
        "        # self.E = 20\n",
        "        self.E = 1\n",
        "        # batch size\n",
        "        self.batch_size = 10\n",
        "        # learning Rate\n",
        "        self.lr = 0.8\n",
        "        # proximal term constant\n",
        "        # self.mu = 0.0\n",
        "        self.mu = 0.001\n",
        "        # percentage of clients to have fewer than E epochs\n",
        "        self.percentage = 0\n",
        "        # self.percentage = 50\n",
        "        # self.percentage = 90\n",
        "        # target test accuracy\n",
        "        self.target_test_accuracy= 99.0\n",
        "        # self.target_test_accuracy=96.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m_JVF83mfM3f"
      },
      "source": [
        "exp_log = dict()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rYOPtnYoqhWd"
      },
      "source": [
        "### IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FRKc7NrzqKpU"
      },
      "source": [
        "train_ds, data_dict, test_ds = iid_partition(corpus, seq_length, val_split=True)  # Not using val_ds but makes train eval periods faster\n",
        "\n",
        "total_clients = len(data_dict.keys())\n",
        "'Total users:', total_clients"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eaKtpKT5q1q_"
      },
      "source": [
        "hparams = Hyperparameters(total_clients)\n",
        "hparams.__dict__"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5ilYMClTV_WR"
      },
      "source": [
        "# Sweeping parameter\n",
        "PARAM_NAME = 'clients_fraction'\n",
        "PARAM_VALUE = hparams.C\n",
        "exp_id = f'{PARAM_NAME}/{PARAM_VALUE}'\n",
        "exp_id"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xAhy4CWVZy3F"
      },
      "source": [
        "EXP_DIR = f'{BASE_DIR}/{exp_id}'\n",
        "os.makedirs(EXP_DIR, exist_ok=True)\n",
        "\n",
        "# tb_logger = SummaryWriter(log_dir)\n",
        "# print(f'TBoard logger created at: {log_dir}')\n",
        "\n",
        "title = 'LSTM FedProx on IID'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LwTdeiv8q8_L"
      },
      "source": [
        "def run_experiment(run_id):\n",
        "\n",
        "    shakespeare_lstm = ShakespeareLSTM(input_dim=seq_length,  \n",
        "                                   embedding_dim=embedding_dim,  \n",
        "                                   hidden_dim=hidden_dim,\n",
        "                                   classes=num_classes,\n",
        "                                   lstm_layers=lstm_layers,\n",
        "                                   dropout=dropout,\n",
        "                                   batch_first=True\n",
        "                                   )\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "        shakespeare_lstm.cuda()\n",
        "    \n",
        "    test_history = []\n",
        "\n",
        "    lstm_iid_trained, test_history = training(shakespeare_lstm,\n",
        "                                            hparams.rounds, hparams.batch_size, hparams.lr,\n",
        "                                            train_ds,\n",
        "                                            data_dict,\n",
        "                                            test_ds,\n",
        "                                            hparams.C, hparams.K, hparams.E, hparams.mu, hparams.percentage,\n",
        "                                            title, \"green\",\n",
        "                                            hparams.target_test_accuracy,\n",
        "                                            corpus, # classes\n",
        "                                            history=test_history,\n",
        "                                            algorithm='fedavg',\n",
        "                                            # tb_logger=tb_writer\n",
        "                                            )\n",
        "    \n",
        "\n",
        "    final_scores = testing(lstm_iid_trained, test_ds, batch_size * 2, nn.CrossEntropyLoss(), len(corpus), corpus)\n",
        "    print(f'\\n\\n========================================================\\n\\n')\n",
        "    print(f'Final scores for Exp {run_id} \\n {final_scores}')\n",
        "\n",
        "    log = {\n",
        "        'history': test_history,\n",
        "        'hyperparams': hparams.__dict__\n",
        "    }\n",
        "\n",
        "    with open(f'{EXP_DIR}/results_iid_{run_id}.pkl', 'wb') as file:\n",
        "        pickle.dump(log, file)\n",
        "\n",
        "    return test_history\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gSU61KsSq87G"
      },
      "source": [
        "exp_history = list()\n",
        "for run_id in range(2):  # TOTAL RUNS\n",
        "    print(f'============== RUNNING EXPERIMENT #{run_id} ==============')\n",
        "    exp_history.append(run_experiment(run_id))\n",
        "    print(f'\\n\\n========================================================\\n\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "us-HifGq3Uhf"
      },
      "source": [
        "exp_log[title] = {\n",
        "    'history': exp_history,\n",
        "    'hyperparams': hparams.__dict__\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qDGpo4ug33dN"
      },
      "source": [
        "df = None\n",
        "for i, e in enumerate(exp_history):\n",
        "    if i == 0:\n",
        "        df = pd.json_normalize(e)\n",
        "        continue\n",
        "    df = df + pd.json_normalize(e)\n",
        "    \n",
        "df_avg = df / len(exp_history)\n",
        "avg_history = df_avg.to_dict(orient='records')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hf77BQAD36Eq"
      },
      "source": [
        "plot_scores(history=avg_history, exp_id=exp_id, title=title, suffix='IID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wJClynRJ38Dh"
      },
      "source": [
        "plot_losses(history=avg_history, exp_id=exp_id, title=title, suffix='IID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N36h40Zs1h_f"
      },
      "source": [
        "with open(f'{EXP_DIR}/results_iid.pkl', 'wb') as file:\n",
        "    pickle.dump(exp_log, file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BaoYWkWgqvUQ"
      },
      "source": [
        "### Non-IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lVmhd_791lk7"
      },
      "source": [
        "exp_log = dict()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pILgaho8qKgF"
      },
      "source": [
        "data_dict, test_ds = noniid_partition(corpus, seq_length=seq_length, val_split=True)\n",
        "\n",
        "total_clients = len(data_dict.keys())\n",
        "'Total users:', total_clients"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y3o7qgBcqKX_"
      },
      "source": [
        "hparams = Hyperparameters(total_clients)\n",
        "hparams.__dict__"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VANr1h0Pq51N"
      },
      "source": [
        "# Sweeping parameter\n",
        "PARAM_NAME = 'clients_fraction'\n",
        "PARAM_VALUE = hparams.C\n",
        "exp_id = f'{PARAM_NAME}/{PARAM_VALUE}'\n",
        "exp_id"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yXgYFIyZ4ipm"
      },
      "source": [
        "EXP_DIR = f'{BASE_DIR}/{exp_id}'\n",
        "os.makedirs(EXP_DIR, exist_ok=True)\n",
        "\n",
        "# tb_logger = SummaryWriter(log_dir)\n",
        "# print(f'TBoard logger created at: {log_dir}')\n",
        "\n",
        "title = 'LSTM FedProx on Non-IID'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vnv7UaE0q6dG"
      },
      "source": [
        "def run_experiment(run_id):\n",
        "\n",
        "    shakespeare_lstm = ShakespeareLSTM(input_dim=seq_length,\n",
        "                                       embedding_dim=embedding_dim,\n",
        "                                       hidden_dim=hidden_dim,\n",
        "                                       classes=num_classes,\n",
        "                                       lstm_layers=lstm_layers,\n",
        "                                       dropout=dropout,\n",
        "                                       batch_first=True\n",
        "                                       )\n",
        "\n",
        "    if torch.cuda.is_available():\n",
        "        shakespeare_lstm.cuda()\n",
        "\n",
        "    test_history = []\n",
        "\n",
        "    lstm_non_iid_trained, test_history = training(shakespeare_lstm,\n",
        "                                                  hparams.rounds, hparams.batch_size, hparams.lr,\n",
        "                                                  None, #  ds empty as it is included in data_dict\n",
        "                                                  data_dict,\n",
        "                                                  test_ds,\n",
        "                                                  hparams.C, hparams.K, hparams.E, hparams.mu, hparams.percentage,\n",
        "                                                  title, \"green\",\n",
        "                                                  hparams.target_test_accuracy,\n",
        "                                                  corpus, # classes\n",
        "                                                  history=test_history,\n",
        "                                                  algorithm='fedavg',\n",
        "                                                  # tb_logger=tb_writer\n",
        "                                                   )\n",
        "\n",
        "    \n",
        "\n",
        "    final_scores = testing(lstm_non_iid_trained, test_ds, batch_size * 2, nn.CrossEntropyLoss(), len(corpus), corpus)\n",
        "    print(f'\\n\\n========================================================\\n\\n')\n",
        "    print(f'Final scores for Exp {run_id} \\n {final_scores}')\n",
        "\n",
        "    log = {\n",
        "        'history': test_history,\n",
        "        'hyperparams': hparams.__dict__\n",
        "    }\n",
        "\n",
        "    with open(f'{EXP_DIR}/results_niid_{run_id}.pkl', 'wb') as file:\n",
        "        pickle.dump(log, file)\n",
        "\n",
        "    return test_history\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0pLbVBwVq6Uw"
      },
      "source": [
        "exp_history = list()\n",
        "for run_id in range(2):  # TOTAL RUNS\n",
        "    print(f'============== RUNNING EXPERIMENT #{run_id} ==============')\n",
        "    exp_history.append(run_experiment(run_id))\n",
        "    print(f'\\n\\n========================================================\\n\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n5F38z5C4qw9"
      },
      "source": [
        "exp_log[title] = {\n",
        "    'history': exp_history,\n",
        "    'hyperparams': hparams.__dict__\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "inIGn3Mh4qpO"
      },
      "source": [
        "df = None\n",
        "for i, e in enumerate(exp_history):\n",
        "    if i == 0:\n",
        "        df = pd.json_normalize(e)\n",
        "        continue\n",
        "    df = df + pd.json_normalize(e)\n",
        "    \n",
        "df_avg = df / len(exp_history)\n",
        "avg_history = df_avg.to_dict(orient='records')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z8ngcls64qjc"
      },
      "source": [
        "plot_scores(history=avg_history, exp_id=exp_id, title=title, suffix='nonIID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GR9vjtYs4qBX"
      },
      "source": [
        "plot_losses(history=avg_history, exp_id=exp_id, title=title, suffix='nonIID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "adK1OTS-40Z8"
      },
      "source": [
        "### Pickle Experiment Results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i5nl-hsa4zqw"
      },
      "source": [
        "with open(f'{EXP_DIR}/results_niid.pkl', 'wb') as file:\n",
        "    pickle.dump(exp_log, file)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}