10 - Federated Learning with flwr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import VisionDataset, Food101\n",
    "from typing import List, Tuple\n",
    "from flwr.common import Metrics\n",
    "\n",
    "import flwr as fl\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import multiprocessing\n",
    "\n",
    "data_path = os.path.join(os.getcwd(),'data', 'food-101')\n",
    "cpu_count = multiprocessing.cpu_count() - 1 # set as you like!\n",
    "#device = torch.device(\"mps\") #CHANGE THIS TO FIT YOUR DEVICE PLEASE :D (maybe under fits)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_size = 4  # number of dataset partions (= number of total clients)\n",
    "\n",
    "client_resources = {\n",
    "        \"num_cpus\": cpu_count\n",
    "}  # each client will get allocated 1 CPUs\n",
    "\n",
    "transformations = transforms.Compose([\n",
    "    transforms.Resize((32,32)),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "\n",
    "# Download Dataset\n",
    "try:\n",
    "    train_data = Food101(data_path, transform=transformations)\n",
    "except:\n",
    "    train_data = Food101(data_path, transform=transformations, download=True) \n",
    "test_data = Food101(data_path, split='test', transform=transformations)\n",
    "\n",
    "lengths = []\n",
    "while sum(lengths) != len(train_data):\n",
    "    lengths = [round(x) for x in np.random.dirichlet(\n",
    "        np.ones(pool_size),size=1)[0] * len(train_data)]\n",
    "    \n",
    "trainloader = DataLoader(train_data, batch_size=32, shuffle=True)\n",
    "testloader = DataLoader(test_data, batch_size=32)\n",
    "num_examples = {\"trainset\" : len(train_data), \"testset\" : len(test_data)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# borrowed from Pytorch quickstart example\n",
    "def train(net, trainloader, epochs, device: str):\n",
    "    \"\"\"Train the network on the training set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss(ignore_index=1)\n",
    "    optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)\n",
    "    net.train()\n",
    "    for _ in range(epochs):\n",
    "        for images, labels in trainloader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            loss = criterion(net(images), labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# borrowed from Pytorch quickstart example\n",
    "def test(net, testloader, device: str):\n",
    "    \"\"\"Validate the network on the entire test set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    correct, total, loss = 0, 0, 0.0\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in testloader:\n",
    "            images, labels = batch[0].to(device), batch[1].to(device)\n",
    "            outputs = net(images)\n",
    "            loss += criterion(outputs, labels).item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    accuracy = correct / total\n",
    "    print(\"Loss: %f, Accuracy: %f\" % (loss, accuracy))\n",
    "    return loss, accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self) -> None:\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 101)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 5 * 5)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "# Load model and data\n",
    "net = Net().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FlowerClient(fl.client.NumPyClient):\n",
    "    \n",
    "    def get_parameters(self, config):\n",
    "        return [val.cpu().numpy() for _, val in net.state_dict().items()]\n",
    "\n",
    "    def set_parameters(self, parameters):\n",
    "        params_dict = zip(net.state_dict().keys(), parameters)\n",
    "        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})\n",
    "        net.load_state_dict(state_dict, strict=True)\n",
    "\n",
    "    def fit(self, parameters, config):\n",
    "        self.set_parameters(parameters)\n",
    "        train(net, trainloader, 3, device)\n",
    "        return self.get_parameters(config={}), num_examples[\"trainset\"], {}\n",
    "\n",
    "    def evaluate(self, parameters, config):\n",
    "        self.set_parameters(parameters)\n",
    "        loss, accuracy = test(net, testloader, device)\n",
    "        return float(loss), num_examples[\"testset\"], {\"accuracy\": float(accuracy)}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "101"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(test_data.classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Before you start training!\n",
    "\n",
    "- Make sure your device is properly set above to fit your compute.\n",
    "- If you have made any changes to this script, download it as a python file and replace the flower/client.py file.\n",
    "- Open a separate terminal and run `python flower/server.py`.\n",
    "- Open 1-3 more terminals and run `python flower/client.py`.\n",
    "- Then run the following cell to also run a client here and watch! :)\n",
    "\n",
    "If you want to change any of the model parameters, structure or even the splits on the data, you'll want to restart the server and clients. Have fun and experiment!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO flwr 2023-02-28 17:38:50,146 | grpc.py:50 | Opened insecure gRPC connection (no certificates were passed)\n",
      "DEBUG flwr 2023-02-28 17:38:50,151 | connection.py:38 | ChannelConnectivity.IDLE\n",
      "DEBUG flwr 2023-02-28 17:38:50,206 | connection.py:38 | ChannelConnectivity.READY\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss: 3458.234933, Accuracy: 0.052713\n",
      "Loss: 3010.837580, Accuracy: 0.136515\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG flwr 2023-02-28 18:44:55,237 | connection.py:109 | gRPC channel closed\n",
      "INFO flwr 2023-02-28 18:44:55,237 | app.py:153 | Disconnect and shut down\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss: 2963.027775, Accuracy: 0.148515\n"
     ]
    }
   ],
   "source": [
    "fl.client.start_numpy_client(\n",
    "    server_address=\"127.0.0.1:8080\",\n",
    "    client=FlowerClient(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Challenges\n",
    "\n",
    "- Adjust the server settings and see how the performance changes (see flower/server.py).\n",
    "- Split the data unevenly across the clients and see how the trianing goes.\n",
    "- Try out another [Flower tutorial](https://flower.dev/docs/quickstart-pytorch.html).\n",
    "- Get a group of several folks together to try running flower in a distributed setup. Document your learnings and share in the reader-contributions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}