federated-cats-vs-dogs-classifier.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Project 3.0.ipynb",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ashwindasr/Federated-Learning/blob/master/federated-cats-vs-dogs-classifier.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MBrS5mM8_buv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "from torchvision import datasets, transforms, models"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "12IH399-FYNo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Uncomment this cell and run it only once.\n",
        "\n",
        "# !git clone https://github.com/OpenMined/PySyft.git\n",
        "# !python ./PySyft/ setup.py test\n",
        "# !pip install syft"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KERONOfGAEyn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import syft as sy                           # Import the Pysyft library                          \n",
        "hook = sy.TorchHook(torch)                  # Hook PyTorch to PySyft\n",
        "bob = sy.VirtualWorker(hook, id=\"bob\")      # Define remote worker bob \n",
        "alice = sy.VirtualWorker(hook, id=\"alice\")  # Define remote worker alice"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yDlNwJT-AK3s",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#Initializing values for arguments\n",
        "class Arguments():\n",
        "    def __init__(self):\n",
        "        self.batch_size = 64\n",
        "        self.test_batch_size = 1000\n",
        "        self.epochs = 20\n",
        "        self.lr = 0.01\n",
        "        self.momentum = 0.5\n",
        "        self.no_cuda = False\n",
        "        self.seed = 1\n",
        "        self.log_interval = 10\n",
        "        self.save_model = False\n",
        "\n",
        "args = Arguments()\n",
        "\n",
        "use_cuda = not args.no_cuda and torch.cuda.is_available()       # Check if GPU is available\n",
        "torch.manual_seed(args.seed)\n",
        "device = torch.device(\"cuda\" if use_cuda else \"cpu\")            # Use set device to CPU or GPU(if available)\n",
        "kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EwD7XlJxnsQS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Run this cell only once\n",
        "\n",
        "!unzip /content/data.zip\n",
        "\n",
        "# Make sure your data folder has the following  folder hierarchy:\n",
        "# \n",
        "#        data\n",
        "#          |- train\n",
        "#          |     |- cats\n",
        "#          |     |- dogs\n",
        "#          |- test\n",
        "#               |- cats\n",
        "#               |- dogs"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QxSFXITgnfmu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Set the path to the extracted data folder\n",
        "data_dir = '/content/data/data'\n",
        "\n",
        "# Define transforms for the training data and testing data\n",
        "train_transforms = transforms.Compose([transforms.RandomRotation(30),\n",
        "                                       transforms.RandomResizedCrop(224),\n",
        "                                       transforms.RandomHorizontalFlip(),\n",
        "                                       transforms.ToTensor(),\n",
        "                                       transforms.Normalize([0.485, 0.456, 0.406],\n",
        "                                                            [0.229, 0.224, 0.225])])\n",
        "\n",
        "test_transforms = transforms.Compose([transforms.Resize(255),\n",
        "                                      transforms.CenterCrop(224),\n",
        "                                      transforms.ToTensor(),\n",
        "                                      transforms.Normalize([0.485, 0.456, 0.406],\n",
        "                                                           [0.229, 0.224, 0.225])])\n",
        "\n",
        "# Load data using Pytorch's ImageFolder dataloader. \n",
        "train_data = datasets.ImageFolder(data_dir + '/train', transform=train_transforms)\n",
        "test_data = datasets.ImageFolder(data_dir + '/test', transform=test_transforms)\n",
        "\n",
        "# Use PySyft's Federated Data Loader to load the federated data.\n",
        "federated_train_loader = sy.FederatedDataLoader( \n",
        "    train_data.federate((bob,alice)), # .federate(()) function distributes the data to the virtual workers\n",
        "    batch_size=64, shuffle=True)\n",
        "\n",
        "# If you encounter an error in the above command, comment out the cell which contains the\n",
        "# installaion files of PySyft, restart runtime and run again.\n",
        "\n",
        "# We use the normal non-federated Dataloader for testing.\n",
        "test_loader = torch.utils.data.DataLoader(test_data, batch_size=64)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bGAkrbr4AVp0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(args, model, device, train_loader, optimizer, epoch):\n",
        "    model.train()\n",
        "    for batch_idx, (data, target) in enumerate(federated_train_loader):\n",
        "        model.send(data.location)         # Send the model to where the data is. (Bob or Alice)\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data)\n",
        "        loss = F.nll_loss(output, target)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        model.get()                       # Get the updated weights from the virtual workers\n",
        "        if batch_idx % args.log_interval == 0:\n",
        "            loss = loss.get()             # Get the new loss from the virtual workers\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                epoch, batch_idx * args.batch_size, len(train_loader) * args.batch_size, \n",
        "                100. * batch_idx / len(train_loader), loss.item()))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cGgmyZ9nAZZx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Test the data in an un-federated way as we normally do. \n",
        "def test(args, model, device, test_loader):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += F.nll_loss(output, target, reduction='sum').item()\n",
        "            pred = output.argmax(1, keepdim=True)  \n",
        "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "\n",
        "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
        "        test_loss, correct, len(test_loader.dataset),\n",
        "        100. * correct / len(test_loader.dataset)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QjDeksJeHaZs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load and modify the model of your convenience, here we use Resnet50\n",
        "model = torchvision.models.resnet50(pretrained=True)\n",
        "\n",
        "model.fc = nn.Sequential(nn.Linear(2048, 512),\n",
        "                                 nn.ReLU(),\n",
        "                                 nn.Dropout(0.2),\n",
        "                                 nn.Linear(512, 10),\n",
        "                                 nn.LogSoftmax(dim=1))  \n",
        "optimizer = optim.SGD(model.parameters(), lr=args.lr)\n",
        "model = model.to(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y3UR4N4cAcvz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train the model and save the weights\n",
        "for epoch in range(1, args.epochs ):\n",
        "    train(args, model, device, federated_train_loader, optimizer, epoch)\n",
        "    test(args, model, device, test_loader)\n",
        "\n",
        "if (args.save_model):\n",
        "    torch.save(model.state_dict(), \"model.pt\")"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}