Shakespeare/qFedAvg.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "qFedAvg.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "x5UTD8mqIDWl",
        "MAtkg7ME8PQn",
        "GpokHKXO8gmS",
        "391xbZuU8vEn",
        "rxAlW1YK9B_b",
        "JtAFxz5h9MRy",
        "QSDdUZjs9W-g",
        "rhN9isJD9Z8f",
        "w8OsHzAt9kDc",
        "p-8XOZlR9wbV",
        "0vQVNPmM93M7",
        "JrTGp6vd-DKa"
      ],
      "toc_visible": true,
      "authorship_tag": "ABX9TyMR+JaDud0PZxZhcegRNqzN",
      "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/qFedAvg.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4v_R7FXY7_fi"
      },
      "source": [
        "# FedPerf - Shakespeare + qFedAvg algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x5UTD8mqIDWl"
      },
      "source": [
        "## Setup & Dependencies Installation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1vTQbtlX731V"
      },
      "source": [
        "%%capture\n",
        "!pip install torchsummaryX unidecode"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FHlNfr2R8EvM"
      },
      "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": "AvWSM6mh8LZr"
      },
      "source": [
        "## Mount GDrive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oh9Lj2-lIIfY"
      },
      "source": [
        "BASE_DIR = '/content/drive/MyDrive/FedPerf/shakespeare/qFedAvg'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "518ez9Oz8LNE"
      },
      "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": "MAtkg7ME8PQn"
      },
      "source": [
        "## Loading Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-6ZYQ2SB8LLo"
      },
      "source": [
        "!rm -Rf data\n",
        "!mkdir -p data scripts"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "skkRzhu98LIx"
      },
      "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": "c9nCIGlB8LGA"
      },
      "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": "jKQNm6cC8LDp"
      },
      "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": "5ewNwmpP8LBg"
      },
      "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": "GpokHKXO8gmS"
      },
      "source": [
        "#### Dataset Preprocessing script"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1YEeiGCT8K--"
      },
      "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": "KE4LT4z48K8u"
      },
      "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": "391xbZuU8vEn"
      },
      "source": [
        "#### Dataset class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5aS5pJenZ8Ow"
      },
      "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": "0r68yAYr8xwJ"
      },
      "source": [
        "##### Federated Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CUTWxbt98K6W"
      },
      "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": "awOBv7tN8_ec"
      },
      "source": [
        "## Partitioning & Data Loaders"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rxAlW1YK9B_b"
      },
      "source": [
        "### IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OqAvrJxm8K4D"
      },
      "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 containing the MNIST Images\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": "0T-NUxsq8K1g"
      },
      "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": "JtAFxz5h9MRy"
      },
      "source": [
        "### Non-IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0w3yiODz8KzT"
      },
      "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": "QSDdUZjs9W-g"
      },
      "source": [
        "## Models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rhN9isJD9Z8f"
      },
      "source": [
        "### Shakespeare LSTM"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GKXxN3HO8Kw3"
      },
      "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": "lwrRH5yD9eZB"
      },
      "source": [
        "#### Model Summary"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zsl5z8CS8Kul"
      },
      "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": "bbJUc6X89jFw"
      },
      "source": [
        "## qFedAvg Algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w8OsHzAt9kDc"
      },
      "source": [
        "### Plot Utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jn1JNhaI8KsX"
      },
      "source": [
        "from sklearn.metrics import f1_score"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oyYu5s4d8KqH"
      },
      "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": "p-8XOZlR9wbV"
      },
      "source": [
        "### Local Training (Client Update)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "83Jn60Gt8KlY"
      },
      "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": "vOn7VSei8KjB"
      },
      "source": [
        "class ClientUpdate(object):\n",
        "  def __init__(self, dataset, batch_size, learning_rate, epochs, idxs, q=None):\n",
        "    \"\"\"\n",
        "\n",
        "    \"\"\"\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.learning_rate = learning_rate\n",
        "    self.epochs = epochs\n",
        "    self.q = q\n",
        "    if not self.q:\n",
        "        # TODO: Client itself adjust fairness \n",
        "        pass\n",
        "    self.mu = 1e-10\n",
        "\n",
        "  def train(self, model):\n",
        "\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.SGD(model.parameters(), lr=self.learning_rate, momentum=0.5)\n",
        "    # optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)\n",
        "\n",
        "    e_loss = []\n",
        "    model_weights = copy.deepcopy(model.state_dict())\n",
        "    for epoch in range(1, self.epochs+1):\n",
        "\n",
        "      train_loss = 0.0\n",
        "\n",
        "      model.train()\n",
        "      # for data, labels in tqdm(self.train_loader):\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",
        "        # print('input', data.size())\n",
        "        output = model(data)\n",
        "        # print('output', output.size())\n",
        "        # print('labels', labels.size())\n",
        "        # calculate the loss\n",
        "        loss = criterion(output, labels)\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",
        "\n",
        "    total_loss = sum(e_loss)/len(e_loss)\n",
        "\n",
        "    # delta weights\n",
        "    model_weights_new = copy.deepcopy(model.state_dict())\n",
        "    L = 1.0 / self.learning_rate\n",
        "\n",
        "    delta_weights, delta, h = {}, {}, {}\n",
        "    loss_q = np.float_power(total_loss + self.mu, self.q)\n",
        "    # updating the global weights\n",
        "    for k in model_weights_new.keys():\n",
        "      delta_weights[k] = (model_weights[k] - model_weights_new[k]) * L\n",
        "      delta[k] =  loss_q * delta_weights[k]\n",
        "      # Estimation of the local Lipchitz constant\n",
        "      h[k] = (self.q * np.float_power(total_loss + self.mu, self.q - 1) * torch.pow(torch.norm(delta_weights[k]), 2)) + (L * loss_q)\n",
        "\n",
        "    return delta, h, total_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0vQVNPmM93M7"
      },
      "source": [
        "### Server Side Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bhOfx0HSCX35"
      },
      "source": [
        "def client_sampling(n, m, weights=None, with_replace=False):\n",
        "    pk = None\n",
        "    if weights:\n",
        "        total_weights = np.sum(np.asarray(weights))\n",
        "        pk = [w * 1.0 / total_weights for w in weights]\n",
        "\n",
        "    return np.random.choice(range(n), m, replace=with_replace, p=pk)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MCJZnkY08Kgs"
      },
      "source": [
        "def training(model, rounds, batch_size, lr, ds, data_dict, test_ds, C, K, E, q=0, sampling='uniform', tb_logger=None, test_history=[], perf_fig_file='loss.jpg'):\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_ds          Dataset used for global testing\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",
        "    - q:               Parameter that tunes the amount of fairness we wish to impose (default: 0 -> vanilla FedAvg objective)\n",
        "    - sampling         Uniform or weighted (default: uniform)\n",
        "    - tb_logger:       Tensorboard SummaryWriter\n",
        "    - test_history:    Test Scores history log\n",
        "    - perf_fig_file    File for storing final performance plot (loss)\n",
        "  Returns:\n",
        "    - model:           Trained model on the server\n",
        "  \"\"\"\n",
        "\n",
        "  # global model weights\n",
        "  global_weights = model.state_dict()\n",
        "\n",
        "  # training loss\n",
        "  train_loss = []\n",
        "\n",
        "  # client weights by total samples\n",
        "  p_k = None\n",
        "  if sampling == 'weighted':\n",
        "    p_k = [len(data_dict[c]) for c in data_dict] if ds else [len(data_dict[c]['train_ds']) for c in data_dict]\n",
        "\n",
        "  # Time log\n",
        "  start_time = time.time()\n",
        "\n",
        "  users_id = list(data_dict.keys())\n",
        "\n",
        "  # Orchestrate training\n",
        "  for curr_round in range(1, rounds+1):\n",
        "    deltas, hs, local_loss = [], [], []\n",
        "\n",
        "    m = max(int(C*K), 1)    \n",
        "    S_t = client_sampling(K, m, weights=p_k, with_replace=False)\n",
        "\n",
        "    print('Round: {} Picking {}/{} clients: {}'.format(curr_round, m, K, S_t))\n",
        "\n",
        "    global_weights = model.state_dict()\n",
        "\n",
        "    for k in tqdm(S_t):\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_, batch_size=batch_size, learning_rate=lr, epochs=E, idxs=idxs, q=q)\n",
        "    #   weights, loss = local_update.train(model=copy.deepcopy(model))\n",
        "      delta_k, h_k, loss = local_update.train(model=copy.deepcopy(model))\n",
        "\n",
        "      deltas.append(copy.deepcopy(delta_k))\n",
        "      hs.append(copy.deepcopy(h_k))\n",
        "      local_loss.append(copy.deepcopy(loss))\n",
        "\n",
        "      if tb_logger:\n",
        "        tb_logger.add_scalar(f'Round/S{k}', loss, curr_round)\n",
        "\n",
        "    # Perform qFedAvg\n",
        "    h_sum = copy.deepcopy(hs[0])\n",
        "    delta_sum = copy.deepcopy(deltas[0])\n",
        "    \n",
        "    for k in h_sum.keys():\n",
        "        for i in range(1, len(hs)):\n",
        "            h_sum[k] += hs[i][k]\n",
        "            delta_sum[k] += deltas[i][k]\n",
        "\n",
        "    new_weights = {}\n",
        "    for k in delta_sum.keys():\n",
        "        for i in range(len(deltas)):\n",
        "            new_weights[k] = delta_sum[k] / h_sum[k]\n",
        "\n",
        "    # Updating global model weights\n",
        "    for k in global_weights.keys():\n",
        "        global_weights[k] -= new_weights[k]\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",
        "\n",
        "    if tb_logger:\n",
        "        tb_logger.add_scalar('Train/Loss', loss_avg, curr_round)\n",
        "        # tb_logger.add_scalar(f'Train/Datapoints', total_datapoints, curr_round)\n",
        "    \n",
        "    # if curr_round % eval_every == 0:\n",
        "    test_scores = testing(model, test_ds, batch_size, nn.CrossEntropyLoss(), num_classes, list(range(num_classes)))\n",
        "    test_scores['train_loss'] = loss_avg\n",
        "    test_history.append(test_scores)\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",
        "  end_time = time.time()\n",
        "  \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, 'tab:orange')\n",
        "\n",
        "  ax.set(xlabel='# Rounds', ylabel='Train Loss',\n",
        "       title=\"Model's Performance with q: {}\".format(q))\n",
        "  ax.grid()\n",
        "  #fig.savefig(perf_fig_file, format='jpg')\n",
        "\n",
        "  print(\"Training Done! Total time: {}\".format(round(end_time - start_time, 3)))\n",
        "  return model, test_history"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JrTGp6vd-DKa"
      },
      "source": [
        "### Testing Loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yEnPyGYb8KeO"
      },
      "source": [
        "def testing(model, dataset, bs, criterion, num_classes, classes, print_all=False):\n",
        "  #test loss \n",
        "  test_loss = 0.0\n",
        "  y_true, y_hat = None, None\n",
        "\n",
        "  correct_class = list(0 for i in range(num_classes))\n",
        "  total_class = list(0 for i in range(num_classes))\n",
        "\n",
        "  if hasattr(dataset, 'dataloader'):\n",
        "    test_loader = dataset.dataloader(batch_size=bs, shuffle=False)\n",
        "  else:\n",
        "    test_loader = DataLoader(dataset, batch_size=bs, shuffle=False)\n",
        "\n",
        "  l = len(test_loader)\n",
        "\n",
        "  model.eval()\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, dim=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",
        "      try:\n",
        "        # print(type(lbl))\n",
        "        # correct_class[lbl.data[0]] += correct.data[i]\n",
        "        correct_class[lbl.item()] += correct[i]\n",
        "        total_class[lbl.item()] += 1\n",
        "      except:\n",
        "          print('Error', lbl, i)\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": "Evpa7UZO-Hii"
      },
      "source": [
        "## Experiments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HhZ2iUQ2am-_"
      },
      "source": [
        "# FAIL-ON-PURPOSE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QXNvhZAw8Kbx"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IadhhL4knbOc"
      },
      "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 = 1  # qFFL\n",
        "        # batch size\n",
        "        self.batch_size = 10  # fedprox\n",
        "        # learning Rate\n",
        "        self.lr = 0.8  # fedprox, qFFL\n",
        "        # fairness\n",
        "        self.q = 0.001  # qFFL\n",
        "        # sampling\n",
        "        # self.sampling = 'uniform'\n",
        "        self.sampling = 'weighted'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VtSTy1GencJ4"
      },
      "source": [
        "exp_log = dict()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hGoeL-HpDv7L"
      },
      "source": [
        "### IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "61pZ-UY6D8nB"
      },
      "source": [
        "train_ds, data_dict, test_ds = iid_partition(corpus, 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": "nUmhoN-zGZp_"
      },
      "source": [
        "hparams = Hyperparameters(total_clients)\n",
        "hparams.__dict__"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XfNgps35XnxY"
      },
      "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": "PmXNlFlYcl2H"
      },
      "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 qFedAvg on IID'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ymcomht5GdBO"
      },
      "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, data_dict, test_ds,\n",
        "                                              hparams.C, hparams.K, hparams.E, hparams.q,\n",
        "                                              sampling=hparams.sampling,\n",
        "                                              test_history=test_history,\n",
        "                                              # tb_logger=tb_logger,\n",
        "                                              # perf_fig_file=f'{BASE_DIR}/loss.jpg'\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": "xhPK_WL9GZhV"
      },
      "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": "a4Rxo_HZn10G"
      },
      "source": [
        "exp_log[title] = {\n",
        "    'history': exp_history,\n",
        "    'hyperparams': hparams.__dict__\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YQGqrZ1_n-e3"
      },
      "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": "H3SO1Cga5tLE"
      },
      "source": [
        "plot_scores(history=avg_history, exp_id=exp_id, title=title, suffix='IID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pwo7EWRh5uQx"
      },
      "source": [
        "plot_losses(history=avg_history, exp_id=exp_id, title=title, suffix='IID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9KPWHl8128He"
      },
      "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": "ttnrJ9FnDxFR"
      },
      "source": [
        "### Non-IID"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9T9hXMvT2_Ud"
      },
      "source": [
        "exp_log = dict()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yQlDBOdJEAaX"
      },
      "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": "f5vIc19G8KUc"
      },
      "source": [
        "hparams = Hyperparameters(total_clients)\n",
        "hparams.__dict__"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WeGGOeN3rFDJ"
      },
      "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": "i3cGeO20ofEZ"
      },
      "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 qFedAvg on Non-IID'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bu1A3GUjGHWy"
      },
      "source": [
        "def run_experiment(run_id):\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_noniid_trained, test_history = training(shakespeare_lstm,\n",
        "                                             hparams.rounds, hparams.batch_size, hparams.lr,\n",
        "                                             None, data_dict, test_ds,\n",
        "                                             hparams.C, hparams.K, hparams.E, hparams.q,\n",
        "                                             sampling=hparams.sampling,\n",
        "                                             test_history=test_history,\n",
        "                                             # tb_logger=tb_logger,\n",
        "                                             # perf_fig_file=f'{BASE_DIR}/loss.jpg'\n",
        "                                             )\n",
        "    \n",
        "    final_scores = testing(lstm_noniid_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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AaaamxrWGKct"
      },
      "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": "ftL-MoHxwe5C"
      },
      "source": [
        "exp_log[title] = {\n",
        "    'history': exp_history,\n",
        "    'hyperparams': hparams.__dict__\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jceXDZXOwezj"
      },
      "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": "1u7uTHwJ6KXE"
      },
      "source": [
        "plot_scores(history=avg_history, exp_id=exp_id, title=title, suffix='nonIID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "um4eBO8O6KLX"
      },
      "source": [
        "plot_losses(history=avg_history, exp_id=exp_id, title=title, suffix='nonIID')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N80BpTFy6aR7"
      },
      "source": [
        "### Pickle Experiment Results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tGlR8COy6aCN"
      },
      "source": [
        "with open(f'{EXP_DIR}/results_niid.pkl', 'wb') as file:\n",
        "    pickle.dump(exp_log, file)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}