{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Launch gRPC server\n",
    "\n",
    "We present how to launch a gRPC server as a federated learning server. Consider only one client so that we can launch a server and a client (from another notebook) together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_clients = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import dependencies\n",
    "\n",
    "We put all the imports here. \n",
    "Our framework `appfl` is backboned by `torch` and its neural network model `torch.nn`. We also import `torchvision` to download the `MNIST` dataset.\n",
    "More importantly, we need to import `appfl.run_grpc_server` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "from appfl.config import Config\n",
    "from appfl.misc.data import Dataset\n",
    "import appfl.run_grpc_server as grpc_server\n",
    "from omegaconf import OmegaConf, DictConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test dataset\n",
    "\n",
    "The server can also hold test data to check the performance of the global model, and the test data needs to be wrapped in `Dataset` object. Note that the server does not need any training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_raw = torchvision.datasets.MNIST(\n",
    "    \"./_data\", train=False, download=False, transform=ToTensor()\n",
    ")\n",
    "test_data_input = []\n",
    "test_data_label = []\n",
    "for idx in range(len(test_data_raw)):\n",
    "    test_data_input.append(test_data_raw[idx][0].tolist())\n",
    "    test_data_label.append(test_data_raw[idx][1])\n",
    "\n",
    "test_dataset = Dataset(\n",
    "    torch.FloatTensor(test_data_input), torch.tensor(test_data_label)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User-defined model\n",
    "\n",
    "Users can define their own models by deriving `torch.nn.Module`. For example in this simulation, we define the following convolutional neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self, num_channel=1, num_classes=10, num_pixel=28):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(\n",
    "            num_channel, 32, kernel_size=5, padding=0, stride=1, bias=True\n",
    "        )\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=0, stride=1, bias=True)\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=(2, 2))\n",
    "        self.act = nn.ReLU(inplace=True)\n",
    "\n",
    "        X = num_pixel\n",
    "        X = math.floor(1 + (X + 2 * 0 - 1 * (5 - 1) - 1) / 1)\n",
    "        X = X / 2\n",
    "        X = math.floor(1 + (X + 2 * 0 - 1 * (5 - 1) - 1) / 1)\n",
    "        X = X / 2\n",
    "        X = int(X)\n",
    "\n",
    "        self.fc1 = nn.Linear(64 * X * X, 512)\n",
    "        self.fc2 = nn.Linear(512, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.act(self.conv1(x))\n",
    "        x = self.maxpool(x)\n",
    "        x = self.act(self.conv2(x))\n",
    "        x = self.maxpool(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.act(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "model = CNN()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User-defined loss and metric\n",
    "We define the loss function and the validation metric for the training as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = torch.nn.CrossEntropyLoss()   \n",
    "\n",
    "def accuracy(y_true, y_pred):\n",
    "    '''\n",
    "    y_true and y_pred are both of type np.ndarray\n",
    "    y_true (N, d) where N is the size of the validation set, and d is the dimension of the label\n",
    "    y_pred (N, D) where N is the size of the validation set, and D is the output dimension of the ML model\n",
    "    '''\n",
    "    if len(y_pred.shape) == 1:\n",
    "        y_pred = np.round(y_pred)\n",
    "    else:\n",
    "        y_pred = y_pred.argmax(axis=1)\n",
    "    return 100*np.sum(y_pred==y_true)/y_pred.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Runs with configuration\n",
    "\n",
    "We run the `appfl` training with the data and model defined above. \n",
    "A number of parameters can be easily set by changing the configuration values.\n",
    "We read the default configurations from `appfl.config.Config` class as a `DictConfig` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fed:\n",
      "  type: federated\n",
      "  servername: ServerFedAvg\n",
      "  clientname: ClientOptim\n",
      "  args:\n",
      "    server_learning_rate: 0.01\n",
      "    server_adapt_param: 0.001\n",
      "    server_momentum_param_1: 0.9\n",
      "    server_momentum_param_2: 0.99\n",
      "    optim: SGD\n",
      "    num_local_epochs: 10\n",
      "    optim_args:\n",
      "      lr: 0.001\n",
      "    use_dp: false\n",
      "    epsilon: 1\n",
      "    clip_grad: false\n",
      "    clip_value: 1\n",
      "    clip_norm: 1\n",
      "device: cpu\n",
      "device_server: cpu\n",
      "num_clients: 1\n",
      "num_epochs: 2\n",
      "num_workers: 0\n",
      "batch_training: true\n",
      "train_data_batch_size: 64\n",
      "train_data_shuffle: true\n",
      "validation: true\n",
      "test_data_batch_size: 64\n",
      "test_data_shuffle: false\n",
      "data_sanity: false\n",
      "reproduce: true\n",
      "pca_dir: ''\n",
      "params_start: 0\n",
      "params_end: 49\n",
      "ncomponents: 40\n",
      "use_tensorboard: false\n",
      "load_model: false\n",
      "load_model_dirname: ''\n",
      "load_model_filename: ''\n",
      "save_model: false\n",
      "save_model_dirname: ''\n",
      "save_model_filename: ''\n",
      "checkpoints_interval: 2\n",
      "save_model_state_dict: false\n",
      "send_final_model: false\n",
      "output_dirname: output\n",
      "output_filename: result\n",
      "logginginfo: {}\n",
      "summary_file: ''\n",
      "personalization: false\n",
      "p_layers: []\n",
      "config_name: ''\n",
      "max_message_size: 104857600\n",
      "operator:\n",
      "  id: 1\n",
      "server:\n",
      "  id: 1\n",
      "  host: localhost\n",
      "  port: 50051\n",
      "  use_tls: false\n",
      "  api_key: null\n",
      "client:\n",
      "  id: 1\n",
      "enable_compression: false\n",
      "lossy_compressor: SZ2\n",
      "lossless_compressor: blosc\n",
      "compressor_sz2_path: ../.compressor/SZ/build/sz/libSZ.dylib\n",
      "compressor_sz3_path: ../.compressor/SZ3/build/tools/sz3c/libSZ3c.dylib\n",
      "compressor_szx_path: ../.compressor/SZx-main/build/lib/libSZx.dylib\n",
      "error_bounding_mode: ''\n",
      "error_bound: 0.0\n",
      "flat_model_dtype: np.float32\n",
      "param_cutoff: 1024\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cfg: DictConfig = OmegaConf.structured(Config)\n",
    "print(OmegaConf.to_yaml(cfg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the server, we just run it by setting the number of global epochs to 5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/zilinghanli/Documents/courses/appfl/APPFL-Dev/src/appfl/comm/grpc/grpc_server.py:189: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/utils/tensor_numpy.cpp:212.)\n",
      "  primal_tensors[name] = torch.from_numpy(nparray)\n",
      "[Round:  001] Finished; all clients have sent their results.\n",
      "[Round:  001] Updating model weights\n",
      "[Round:  001] Test set: Average loss: 0.3030, Accuracy: 91.30%, Best Accuracy: 91.30%\n",
      "[Round:  002] Finished; all clients have sent their results.\n",
      "[Round:  002] Updating model weights\n",
      "[Round:  002] Test set: Average loss: 0.1731, Accuracy: 94.72%, Best Accuracy: 94.72%\n",
      "[Round:  003] Finished; all clients have sent their results.\n",
      "[Round:  003] Updating model weights\n",
      "[Round:  003] Test set: Average loss: 0.1242, Accuracy: 96.51%, Best Accuracy: 96.51%\n",
      "[Round:  004] Finished; all clients have sent their results.\n",
      "[Round:  004] Updating model weights\n",
      "[Round:  004] Test set: Average loss: 0.0844, Accuracy: 97.53%, Best Accuracy: 97.53%\n",
      "[Round:  005] Finished; all clients have sent their results.\n",
      "[Round:  005] Updating model weights\n",
      "[Round:  005] Test set: Average loss: 0.0758, Accuracy: 97.56%, Best Accuracy: 97.56%\n"
     ]
    }
   ],
   "source": [
    "cfg.num_epochs = 5\n",
    "grpc_server.run_server(cfg, model, loss_fn, num_clients, test_dataset, accuracy)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d5a3775820edfef7d27663833b7a57b274657051daef716a62aaac9a7002010d"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 64-bit ('appfl-dev': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}