Federated Learning - MNIST.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Federated Learning - MNIST.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/resilientmax/60DaysOfUdacity/blob/master/Federated%20Learning%20-%20MNIST.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dv6BtROFtQJ6",
        "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",
        "from torchvision import datasets, transforms"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Q7iqboItWJ2",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "23f838c6-bb89-4efb-f718-ed555696fb57"
      },
      "source": [
        "!pip install syft"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting syft\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/38/2e/16bdefc78eb089e1efa9704c33b8f76f035a30dc935bedd7cbb22f6dabaa/syft-0.1.21a1-py3-none-any.whl (219kB)\n",
            "\u001b[K     |████████████████████████████████| 225kB 4.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: tblib>=1.4.0 in /usr/local/lib/python3.6/dist-packages (from syft) (1.4.0)\n",
            "Requirement already satisfied: torch>=1.1 in /usr/local/lib/python3.6/dist-packages (from syft) (1.1.0)\n",
            "Requirement already satisfied: torchvision>=0.3.0 in /usr/local/lib/python3.6/dist-packages (from syft) (0.3.0)\n",
            "Collecting lz4>=2.1.6 (from syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/0a/c6/96bbb3525a63ebc53ea700cc7d37ab9045542d33b4d262d0f0408ad9bbf2/lz4-2.1.10-cp36-cp36m-manylinux1_x86_64.whl (385kB)\n",
            "\u001b[K     |████████████████████████████████| 389kB 41.1MB/s \n",
            "\u001b[?25hCollecting websocket-client>=0.56.0 (from syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/29/19/44753eab1fdb50770ac69605527e8859468f3c0fd7dc5a76dd9c4dbd7906/websocket_client-0.56.0-py2.py3-none-any.whl (200kB)\n",
            "\u001b[K     |████████████████████████████████| 204kB 41.3MB/s \n",
            "\u001b[?25hCollecting msgpack>=0.6.1 (from syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/92/7e/ae9e91c1bb8d846efafd1f353476e3fd7309778b582d2fb4cea4cc15b9a2/msgpack-0.6.1-cp36-cp36m-manylinux1_x86_64.whl (248kB)\n",
            "\u001b[K     |████████████████████████████████| 256kB 35.8MB/s \n",
            "\u001b[?25hCollecting websockets>=7.0 (from syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/c1/d2/bf72435a7d56f94b57efdeae26c76bf0d16f409fd44ff595da745c3fbefd/websockets-8.0.1-cp36-cp36m-manylinux1_x86_64.whl (72kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 24.0MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.14.0 in /usr/local/lib/python3.6/dist-packages (from syft) (1.16.4)\n",
            "Collecting tf-encrypted>=0.5.4 (from syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/55/ff/7dbd5fc77fcec0df1798268a6b72a2ab0150b854761bc39c77d566798f0b/tf_encrypted-0.5.7-py3-none-manylinux1_x86_64.whl (2.1MB)\n",
            "\u001b[K     |████████████████████████████████| 2.1MB 37.7MB/s \n",
            "\u001b[?25hCollecting zstd>=1.4.0.0 (from syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/22/37/6a7ba746ebddbd6cd06de84367515d6bc239acd94fb3e0b1c85788176ca2/zstd-1.4.1.0.tar.gz (454kB)\n",
            "\u001b[K     |████████████████████████████████| 460kB 39.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: scikit-learn>=0.21.0 in /usr/local/lib/python3.6/dist-packages (from syft) (0.21.2)\n",
            "Collecting flask-socketio>=3.3.2 (from syft)\n",
            "  Downloading https://files.pythonhosted.org/packages/4b/68/fe4806d3a0a5909d274367eb9b3b87262906c1515024f46c2443a36a0c82/Flask_SocketIO-4.1.0-py2.py3-none-any.whl\n",
            "Requirement already satisfied: Flask>=1.0.2 in /usr/local/lib/python3.6/dist-packages (from syft) (1.1.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision>=0.3.0->syft) (1.12.0)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision>=0.3.0->syft) (4.3.0)\n",
            "Requirement already satisfied: tensorflow<2,>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tf-encrypted>=0.5.4->syft) (1.14.0)\n",
            "Collecting pyyaml>=5.1 (from tf-encrypted>=0.5.4->syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a3/65/837fefac7475963d1eccf4aa684c23b95aa6c1d033a2c5965ccb11e22623/PyYAML-5.1.1.tar.gz (274kB)\n",
            "\u001b[K     |████████████████████████████████| 276kB 39.3MB/s \n",
            "\u001b[?25hRequirement already satisfied: scipy>=0.17.0 in /usr/local/lib/python3.6/dist-packages (from scikit-learn>=0.21.0->syft) (1.3.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn>=0.21.0->syft) (0.13.2)\n",
            "Collecting python-socketio>=2.1.0 (from flask-socketio>=3.3.2->syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/26/1b/57e860a86f2a01be86ae1dacfa0cd8c4dfbfcd4593322268b61b5a07b564/python_socketio-4.2.0-py2.py3-none-any.whl (46kB)\n",
            "\u001b[K     |████████████████████████████████| 51kB 18.0MB/s \n",
            "\u001b[?25hRequirement already satisfied: click>=5.1 in /usr/local/lib/python3.6/dist-packages (from Flask>=1.0.2->syft) (7.0)\n",
            "Requirement already satisfied: Jinja2>=2.10.1 in /usr/local/lib/python3.6/dist-packages (from Flask>=1.0.2->syft) (2.10.1)\n",
            "Requirement already satisfied: Werkzeug>=0.15 in /usr/local/lib/python3.6/dist-packages (from Flask>=1.0.2->syft) (0.15.5)\n",
            "Requirement already satisfied: itsdangerous>=0.24 in /usr/local/lib/python3.6/dist-packages (from Flask>=1.0.2->syft) (1.1.0)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision>=0.3.0->syft) (0.46)\n",
            "Requirement already satisfied: tensorboard<1.15.0,>=1.14.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.14.0)\n",
            "Requirement already satisfied: tensorflow-estimator<1.15.0rc0,>=1.14.0rc0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.14.0)\n",
            "Requirement already satisfied: google-pasta>=0.1.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (0.1.7)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (0.7.1)\n",
            "Requirement already satisfied: keras-applications>=1.0.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.0.8)\n",
            "Requirement already satisfied: astor>=0.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (0.8.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.1.0)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.11.2)\n",
            "Requirement already satisfied: protobuf>=3.6.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (3.7.1)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.15.0)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (0.33.4)\n",
            "Requirement already satisfied: gast>=0.2.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (0.2.2)\n",
            "Requirement already satisfied: keras-preprocessing>=1.0.5 in /usr/local/lib/python3.6/dist-packages (from tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (1.1.0)\n",
            "Collecting python-engineio>=3.8.0 (from python-socketio>=2.1.0->flask-socketio>=3.3.2->syft)\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/bd/b8/0fc389ca5c445051b37b17802f80bbf1b51c1e3b48b772ee608efbb90583/python_engineio-3.8.2.post1-py2.py3-none-any.whl (119kB)\n",
            "\u001b[K     |████████████████████████████████| 122kB 36.7MB/s \n",
            "\u001b[?25hRequirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.6/dist-packages (from Jinja2>=2.10.1->Flask>=1.0.2->syft) (1.1.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tensorboard<1.15.0,>=1.14.0->tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (3.1.1)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<1.15.0,>=1.14.0->tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (41.0.1)\n",
            "Requirement already satisfied: h5py in /usr/local/lib/python3.6/dist-packages (from keras-applications>=1.0.6->tensorflow<2,>=1.12.0->tf-encrypted>=0.5.4->syft) (2.8.0)\n",
            "Building wheels for collected packages: zstd, pyyaml\n",
            "  Building wheel for zstd (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Stored in directory: /root/.cache/pip/wheels/66/3f/ee/ac08c81af7c1b24a80c746df669ea3cb37542d27877d66ccf4\n",
            "  Building wheel for pyyaml (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Stored in directory: /root/.cache/pip/wheels/16/27/a1/775c62ddea7bfa62324fd1f65847ed31c55dadb6051481ba3f\n",
            "Successfully built zstd pyyaml\n",
            "Installing collected packages: lz4, websocket-client, msgpack, websockets, pyyaml, tf-encrypted, zstd, python-engineio, python-socketio, flask-socketio, syft\n",
            "  Found existing installation: msgpack 0.5.6\n",
            "    Uninstalling msgpack-0.5.6:\n",
            "      Successfully uninstalled msgpack-0.5.6\n",
            "  Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "Successfully installed flask-socketio-4.1.0 lz4-2.1.10 msgpack-0.6.1 python-engineio-3.8.2.post1 python-socketio-4.2.0 pyyaml-5.1.1 syft-0.1.21a1 tf-encrypted-0.5.7 websocket-client-0.56.0 websockets-8.0.1 zstd-1.4.1.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I-3C2q4utZKl",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 106
        },
        "outputId": "a20f8409-c0b7-461a-d996-194ee61fc35a"
      },
      "source": [
        "import syft as sy  # <-- NEW: import the Pysyft library\n",
        "hook = sy.TorchHook(torch)  # <-- NEW: hook PyTorch ie add extra functionalities to support Federated Learning\n",
        "bob = sy.VirtualWorker(hook, id=\"bob\")  # <-- NEW: define remote worker bob\n",
        "alice = sy.VirtualWorker(hook, id=\"alice\")  # <-- NEW: and alice"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING: Logging before flag parsing goes to stderr.\n",
            "W0722 12:56:57.738648 140124902066048 secure_random.py:26] Falling back to insecure randomness since the required custom op could not be found for the installed version of TensorFlow. Fix this by compiling custom ops. Missing file was '/usr/local/lib/python3.6/dist-packages/tf_encrypted/operations/secure_random/secure_random_module_tf_1.14.0.so'\n",
            "W0722 12:56:57.757672 140124902066048 deprecation_wrapper.py:119] From /usr/local/lib/python3.6/dist-packages/tf_encrypted/session.py:26: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.\n",
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d9s9Zxnutde2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Arguments():\n",
        "    def __init__(self):\n",
        "        self.batch_size = 64\n",
        "        self.test_batch_size = 1000\n",
        "        self.epochs = 2\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()\n",
        "\n",
        "torch.manual_seed(args.seed)\n",
        "\n",
        "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
        "\n",
        "kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2rI02Xgrtf8d",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 260
        },
        "outputId": "c0eb9d43-d7aa-4334-f39b-eff54a74531e"
      },
      "source": [
        "federated_train_loader = sy.FederatedDataLoader( # <-- this is now a FederatedDataLoader \n",
        "    datasets.MNIST('../data', train=True, download=True,\n",
        "                   transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ]))\n",
        "    .federate((bob, alice)), # <-- NEW: we distribute the dataset across all the workers, it's now a FederatedDataset\n",
        "    batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('../data', train=False, transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ])),\n",
        "    batch_size=args.test_batch_size, shuffle=True, **kwargs)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\r0it [00:00, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ../data/MNIST/raw/train-images-idx3-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "9920512it [00:01, 9137977.43it/s]                            \n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Extracting ../data/MNIST/raw/train-images-idx3-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "  0%|          | 0/28881 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ../data/MNIST/raw/train-labels-idx1-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "32768it [00:00, 135012.64it/s]           \n",
            "  0%|          | 0/1648877 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Extracting ../data/MNIST/raw/train-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "1654784it [00:00, 2308564.60it/s]                            \n",
            "0it [00:00, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Extracting ../data/MNIST/raw/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "8192it [00:00, 54078.14it/s]            \n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Extracting ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n",
            "Processing...\n",
            "Done!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zeDHOcqjtj1Y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 20, 5, 1)\n",
        "        self.conv2 = nn.Conv2d(20, 50, 5, 1)\n",
        "        self.fc1 = nn.Linear(4*4*50, 500)\n",
        "        self.fc2 = nn.Linear(500, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(self.conv1(x))\n",
        "        x = F.max_pool2d(x, 2, 2)\n",
        "        x = F.relu(self.conv2(x))\n",
        "        x = F.max_pool2d(x, 2, 2)\n",
        "        x = x.view(-1, 4*4*50)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        return F.log_softmax(x, dim=1)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CRG00A3qtoAn",
        "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): # <-- now it is a distributed dataset\n",
        "        model.send(data.location) # <-- NEW: send the model to the right location\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() # <-- NEW: get the model back\n",
        "        if batch_idx % args.log_interval == 0:\n",
        "            loss = loss.get() # <-- NEW: get the loss back\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                epoch, batch_idx * args.batch_size, len(train_loader) * args.batch_size, #batch_idx * len(data), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), loss.item()))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7BD4aGZbtq3B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "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() # sum up batch loss\n",
        "            pred = output.argmax(1, keepdim=True) # get the index of the max log-probability \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)))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZdEkWwzRtt-c",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "d50d5fa3-6657-44a0-de46-d1cebeb07b5a"
      },
      "source": [
        "model = Net().to(device)\n",
        "optimizer = optim.SGD(model.parameters(), lr=args.lr) \n",
        "for epoch in range(1, args.epochs + 1):\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(), \"mnist_cnn.pt\")"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 1 [0/60032 (0%)]\tLoss: 2.305134\n",
            "Train Epoch: 1 [640/60032 (1%)]\tLoss: 2.273475\n",
            "Train Epoch: 1 [1280/60032 (2%)]\tLoss: 2.216174\n",
            "Train Epoch: 1 [1920/60032 (3%)]\tLoss: 2.156802\n",
            "Train Epoch: 1 [2560/60032 (4%)]\tLoss: 2.139429\n",
            "Train Epoch: 1 [3200/60032 (5%)]\tLoss: 2.053059\n",
            "Train Epoch: 1 [3840/60032 (6%)]\tLoss: 1.896587\n",
            "Train Epoch: 1 [4480/60032 (7%)]\tLoss: 1.917239\n",
            "Train Epoch: 1 [5120/60032 (9%)]\tLoss: 1.655076\n",
            "Train Epoch: 1 [5760/60032 (10%)]\tLoss: 1.440329\n",
            "Train Epoch: 1 [6400/60032 (11%)]\tLoss: 1.231347\n",
            "Train Epoch: 1 [7040/60032 (12%)]\tLoss: 0.983715\n",
            "Train Epoch: 1 [7680/60032 (13%)]\tLoss: 0.867023\n",
            "Train Epoch: 1 [8320/60032 (14%)]\tLoss: 0.890953\n",
            "Train Epoch: 1 [8960/60032 (15%)]\tLoss: 0.861902\n",
            "Train Epoch: 1 [9600/60032 (16%)]\tLoss: 0.654317\n",
            "Train Epoch: 1 [10240/60032 (17%)]\tLoss: 0.587091\n",
            "Train Epoch: 1 [10880/60032 (18%)]\tLoss: 0.693011\n",
            "Train Epoch: 1 [11520/60032 (19%)]\tLoss: 0.593099\n",
            "Train Epoch: 1 [12160/60032 (20%)]\tLoss: 0.531986\n",
            "Train Epoch: 1 [12800/60032 (21%)]\tLoss: 0.400264\n",
            "Train Epoch: 1 [13440/60032 (22%)]\tLoss: 0.455322\n",
            "Train Epoch: 1 [14080/60032 (23%)]\tLoss: 0.438987\n",
            "Train Epoch: 1 [14720/60032 (25%)]\tLoss: 0.398361\n",
            "Train Epoch: 1 [15360/60032 (26%)]\tLoss: 0.371324\n",
            "Train Epoch: 1 [16000/60032 (27%)]\tLoss: 0.289150\n",
            "Train Epoch: 1 [16640/60032 (28%)]\tLoss: 0.416670\n",
            "Train Epoch: 1 [17280/60032 (29%)]\tLoss: 0.304617\n",
            "Train Epoch: 1 [17920/60032 (30%)]\tLoss: 0.367789\n",
            "Train Epoch: 1 [18560/60032 (31%)]\tLoss: 0.386265\n",
            "Train Epoch: 1 [19200/60032 (32%)]\tLoss: 0.314097\n",
            "Train Epoch: 1 [19840/60032 (33%)]\tLoss: 0.238745\n",
            "Train Epoch: 1 [20480/60032 (34%)]\tLoss: 0.534601\n",
            "Train Epoch: 1 [21120/60032 (35%)]\tLoss: 0.369211\n",
            "Train Epoch: 1 [21760/60032 (36%)]\tLoss: 0.464785\n",
            "Train Epoch: 1 [22400/60032 (37%)]\tLoss: 0.279363\n",
            "Train Epoch: 1 [23040/60032 (38%)]\tLoss: 0.238020\n",
            "Train Epoch: 1 [23680/60032 (39%)]\tLoss: 0.182981\n",
            "Train Epoch: 1 [24320/60032 (41%)]\tLoss: 0.321955\n",
            "Train Epoch: 1 [24960/60032 (42%)]\tLoss: 0.187752\n",
            "Train Epoch: 1 [25600/60032 (43%)]\tLoss: 0.286318\n",
            "Train Epoch: 1 [26240/60032 (44%)]\tLoss: 0.365828\n",
            "Train Epoch: 1 [26880/60032 (45%)]\tLoss: 0.523223\n",
            "Train Epoch: 1 [27520/60032 (46%)]\tLoss: 0.147624\n",
            "Train Epoch: 1 [28160/60032 (47%)]\tLoss: 0.131381\n",
            "Train Epoch: 1 [28800/60032 (48%)]\tLoss: 0.224323\n",
            "Train Epoch: 1 [29440/60032 (49%)]\tLoss: 0.274111\n",
            "Train Epoch: 1 [30080/60032 (50%)]\tLoss: 0.215077\n",
            "Train Epoch: 1 [30720/60032 (51%)]\tLoss: 0.143916\n",
            "Train Epoch: 1 [31360/60032 (52%)]\tLoss: 0.291434\n",
            "Train Epoch: 1 [32000/60032 (53%)]\tLoss: 0.230507\n",
            "Train Epoch: 1 [32640/60032 (54%)]\tLoss: 0.268499\n",
            "Train Epoch: 1 [33280/60032 (55%)]\tLoss: 0.205015\n",
            "Train Epoch: 1 [33920/60032 (57%)]\tLoss: 0.147734\n",
            "Train Epoch: 1 [34560/60032 (58%)]\tLoss: 0.187934\n",
            "Train Epoch: 1 [35200/60032 (59%)]\tLoss: 0.171631\n",
            "Train Epoch: 1 [35840/60032 (60%)]\tLoss: 0.478283\n",
            "Train Epoch: 1 [36480/60032 (61%)]\tLoss: 0.303621\n",
            "Train Epoch: 1 [37120/60032 (62%)]\tLoss: 0.243553\n",
            "Train Epoch: 1 [37760/60032 (63%)]\tLoss: 0.263034\n",
            "Train Epoch: 1 [38400/60032 (64%)]\tLoss: 0.239960\n",
            "Train Epoch: 1 [39040/60032 (65%)]\tLoss: 0.456759\n",
            "Train Epoch: 1 [39680/60032 (66%)]\tLoss: 0.200484\n",
            "Train Epoch: 1 [40320/60032 (67%)]\tLoss: 0.256438\n",
            "Train Epoch: 1 [40960/60032 (68%)]\tLoss: 0.196973\n",
            "Train Epoch: 1 [41600/60032 (69%)]\tLoss: 0.270759\n",
            "Train Epoch: 1 [42240/60032 (70%)]\tLoss: 0.191681\n",
            "Train Epoch: 1 [42880/60032 (71%)]\tLoss: 0.118813\n",
            "Train Epoch: 1 [43520/60032 (72%)]\tLoss: 0.176458\n",
            "Train Epoch: 1 [44160/60032 (74%)]\tLoss: 0.174667\n",
            "Train Epoch: 1 [44800/60032 (75%)]\tLoss: 0.193119\n",
            "Train Epoch: 1 [45440/60032 (76%)]\tLoss: 0.349294\n",
            "Train Epoch: 1 [46080/60032 (77%)]\tLoss: 0.220412\n",
            "Train Epoch: 1 [46720/60032 (78%)]\tLoss: 0.191737\n",
            "Train Epoch: 1 [47360/60032 (79%)]\tLoss: 0.155758\n",
            "Train Epoch: 1 [48000/60032 (80%)]\tLoss: 0.323814\n",
            "Train Epoch: 1 [48640/60032 (81%)]\tLoss: 0.246998\n",
            "Train Epoch: 1 [49280/60032 (82%)]\tLoss: 0.254214\n",
            "Train Epoch: 1 [49920/60032 (83%)]\tLoss: 0.273938\n",
            "Train Epoch: 1 [50560/60032 (84%)]\tLoss: 0.266370\n",
            "Train Epoch: 1 [51200/60032 (85%)]\tLoss: 0.150609\n",
            "Train Epoch: 1 [51840/60032 (86%)]\tLoss: 0.129327\n",
            "Train Epoch: 1 [52480/60032 (87%)]\tLoss: 0.276623\n",
            "Train Epoch: 1 [53120/60032 (88%)]\tLoss: 0.212820\n",
            "Train Epoch: 1 [53760/60032 (90%)]\tLoss: 0.182833\n",
            "Train Epoch: 1 [54400/60032 (91%)]\tLoss: 0.143227\n",
            "Train Epoch: 1 [55040/60032 (92%)]\tLoss: 0.118357\n",
            "Train Epoch: 1 [55680/60032 (93%)]\tLoss: 0.223088\n",
            "Train Epoch: 1 [56320/60032 (94%)]\tLoss: 0.214281\n",
            "Train Epoch: 1 [56960/60032 (95%)]\tLoss: 0.085184\n",
            "Train Epoch: 1 [57600/60032 (96%)]\tLoss: 0.081163\n",
            "Train Epoch: 1 [58240/60032 (97%)]\tLoss: 0.119667\n",
            "Train Epoch: 1 [58880/60032 (98%)]\tLoss: 0.172809\n",
            "Train Epoch: 1 [59520/60032 (99%)]\tLoss: 0.143191\n",
            "\n",
            "Test set: Average loss: 0.1575, Accuracy: 9511/10000 (95%)\n",
            "\n",
            "Train Epoch: 2 [0/60032 (0%)]\tLoss: 0.103221\n",
            "Train Epoch: 2 [640/60032 (1%)]\tLoss: 0.244422\n",
            "Train Epoch: 2 [1280/60032 (2%)]\tLoss: 0.402155\n",
            "Train Epoch: 2 [1920/60032 (3%)]\tLoss: 0.105845\n",
            "Train Epoch: 2 [2560/60032 (4%)]\tLoss: 0.349342\n",
            "Train Epoch: 2 [3200/60032 (5%)]\tLoss: 0.211654\n",
            "Train Epoch: 2 [3840/60032 (6%)]\tLoss: 0.147023\n",
            "Train Epoch: 2 [4480/60032 (7%)]\tLoss: 0.192190\n",
            "Train Epoch: 2 [5120/60032 (9%)]\tLoss: 0.103523\n",
            "Train Epoch: 2 [5760/60032 (10%)]\tLoss: 0.148321\n",
            "Train Epoch: 2 [6400/60032 (11%)]\tLoss: 0.148263\n",
            "Train Epoch: 2 [7040/60032 (12%)]\tLoss: 0.060192\n",
            "Train Epoch: 2 [7680/60032 (13%)]\tLoss: 0.108846\n",
            "Train Epoch: 2 [8320/60032 (14%)]\tLoss: 0.152352\n",
            "Train Epoch: 2 [8960/60032 (15%)]\tLoss: 0.152677\n",
            "Train Epoch: 2 [9600/60032 (16%)]\tLoss: 0.111557\n",
            "Train Epoch: 2 [10240/60032 (17%)]\tLoss: 0.186860\n",
            "Train Epoch: 2 [10880/60032 (18%)]\tLoss: 0.138993\n",
            "Train Epoch: 2 [11520/60032 (19%)]\tLoss: 0.118514\n",
            "Train Epoch: 2 [12160/60032 (20%)]\tLoss: 0.221352\n",
            "Train Epoch: 2 [12800/60032 (21%)]\tLoss: 0.132464\n",
            "Train Epoch: 2 [13440/60032 (22%)]\tLoss: 0.063139\n",
            "Train Epoch: 2 [14080/60032 (23%)]\tLoss: 0.124831\n",
            "Train Epoch: 2 [14720/60032 (25%)]\tLoss: 0.163303\n",
            "Train Epoch: 2 [15360/60032 (26%)]\tLoss: 0.089060\n",
            "Train Epoch: 2 [16000/60032 (27%)]\tLoss: 0.148574\n",
            "Train Epoch: 2 [16640/60032 (28%)]\tLoss: 0.046008\n",
            "Train Epoch: 2 [17280/60032 (29%)]\tLoss: 0.156941\n",
            "Train Epoch: 2 [17920/60032 (30%)]\tLoss: 0.132319\n",
            "Train Epoch: 2 [18560/60032 (31%)]\tLoss: 0.099179\n",
            "Train Epoch: 2 [19200/60032 (32%)]\tLoss: 0.159745\n",
            "Train Epoch: 2 [19840/60032 (33%)]\tLoss: 0.054083\n",
            "Train Epoch: 2 [20480/60032 (34%)]\tLoss: 0.210852\n",
            "Train Epoch: 2 [21120/60032 (35%)]\tLoss: 0.157456\n",
            "Train Epoch: 2 [21760/60032 (36%)]\tLoss: 0.115076\n",
            "Train Epoch: 2 [22400/60032 (37%)]\tLoss: 0.144074\n",
            "Train Epoch: 2 [23040/60032 (38%)]\tLoss: 0.229732\n",
            "Train Epoch: 2 [23680/60032 (39%)]\tLoss: 0.074357\n",
            "Train Epoch: 2 [24320/60032 (41%)]\tLoss: 0.181132\n",
            "Train Epoch: 2 [24960/60032 (42%)]\tLoss: 0.197346\n",
            "Train Epoch: 2 [25600/60032 (43%)]\tLoss: 0.067677\n",
            "Train Epoch: 2 [26240/60032 (44%)]\tLoss: 0.065908\n",
            "Train Epoch: 2 [26880/60032 (45%)]\tLoss: 0.206655\n",
            "Train Epoch: 2 [27520/60032 (46%)]\tLoss: 0.115207\n",
            "Train Epoch: 2 [28160/60032 (47%)]\tLoss: 0.153855\n",
            "Train Epoch: 2 [28800/60032 (48%)]\tLoss: 0.079507\n",
            "Train Epoch: 2 [29440/60032 (49%)]\tLoss: 0.126364\n",
            "Train Epoch: 2 [30080/60032 (50%)]\tLoss: 0.038445\n",
            "Train Epoch: 2 [30720/60032 (51%)]\tLoss: 0.063148\n",
            "Train Epoch: 2 [31360/60032 (52%)]\tLoss: 0.074466\n",
            "Train Epoch: 2 [32000/60032 (53%)]\tLoss: 0.056466\n",
            "Train Epoch: 2 [32640/60032 (54%)]\tLoss: 0.158469\n",
            "Train Epoch: 2 [33280/60032 (55%)]\tLoss: 0.163727\n",
            "Train Epoch: 2 [33920/60032 (57%)]\tLoss: 0.109543\n",
            "Train Epoch: 2 [34560/60032 (58%)]\tLoss: 0.156446\n",
            "Train Epoch: 2 [35200/60032 (59%)]\tLoss: 0.072912\n",
            "Train Epoch: 2 [35840/60032 (60%)]\tLoss: 0.130330\n",
            "Train Epoch: 2 [36480/60032 (61%)]\tLoss: 0.074369\n",
            "Train Epoch: 2 [37120/60032 (62%)]\tLoss: 0.093282\n",
            "Train Epoch: 2 [37760/60032 (63%)]\tLoss: 0.090162\n",
            "Train Epoch: 2 [38400/60032 (64%)]\tLoss: 0.161415\n",
            "Train Epoch: 2 [39040/60032 (65%)]\tLoss: 0.128906\n",
            "Train Epoch: 2 [39680/60032 (66%)]\tLoss: 0.049100\n",
            "Train Epoch: 2 [40320/60032 (67%)]\tLoss: 0.074256\n",
            "Train Epoch: 2 [40960/60032 (68%)]\tLoss: 0.113933\n",
            "Train Epoch: 2 [41600/60032 (69%)]\tLoss: 0.122479\n",
            "Train Epoch: 2 [42240/60032 (70%)]\tLoss: 0.153609\n",
            "Train Epoch: 2 [42880/60032 (71%)]\tLoss: 0.138244\n",
            "Train Epoch: 2 [43520/60032 (72%)]\tLoss: 0.107693\n",
            "Train Epoch: 2 [44160/60032 (74%)]\tLoss: 0.048192\n",
            "Train Epoch: 2 [44800/60032 (75%)]\tLoss: 0.102147\n",
            "Train Epoch: 2 [45440/60032 (76%)]\tLoss: 0.116528\n",
            "Train Epoch: 2 [46080/60032 (77%)]\tLoss: 0.085225\n",
            "Train Epoch: 2 [46720/60032 (78%)]\tLoss: 0.072129\n",
            "Train Epoch: 2 [47360/60032 (79%)]\tLoss: 0.198065\n",
            "Train Epoch: 2 [48000/60032 (80%)]\tLoss: 0.101033\n",
            "Train Epoch: 2 [48640/60032 (81%)]\tLoss: 0.082122\n",
            "Train Epoch: 2 [49280/60032 (82%)]\tLoss: 0.095091\n",
            "Train Epoch: 2 [49920/60032 (83%)]\tLoss: 0.154688\n",
            "Train Epoch: 2 [50560/60032 (84%)]\tLoss: 0.106988\n",
            "Train Epoch: 2 [51200/60032 (85%)]\tLoss: 0.062860\n",
            "Train Epoch: 2 [51840/60032 (86%)]\tLoss: 0.032838\n",
            "Train Epoch: 2 [52480/60032 (87%)]\tLoss: 0.167880\n",
            "Train Epoch: 2 [53120/60032 (88%)]\tLoss: 0.069700\n",
            "Train Epoch: 2 [53760/60032 (90%)]\tLoss: 0.073842\n",
            "Train Epoch: 2 [54400/60032 (91%)]\tLoss: 0.049357\n",
            "Train Epoch: 2 [55040/60032 (92%)]\tLoss: 0.098377\n",
            "Train Epoch: 2 [55680/60032 (93%)]\tLoss: 0.113492\n",
            "Train Epoch: 2 [56320/60032 (94%)]\tLoss: 0.091218\n",
            "Train Epoch: 2 [56960/60032 (95%)]\tLoss: 0.041372\n",
            "Train Epoch: 2 [57600/60032 (96%)]\tLoss: 0.111734\n",
            "Train Epoch: 2 [58240/60032 (97%)]\tLoss: 0.083618\n",
            "Train Epoch: 2 [58880/60032 (98%)]\tLoss: 0.064205\n",
            "Train Epoch: 2 [59520/60032 (99%)]\tLoss: 0.069066\n",
            "\n",
            "Test set: Average loss: 0.0899, Accuracy: 9735/10000 (97%)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w4V-iqPZt8Ry",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}