FedGAN-HE-DP-MNIST-Server.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f265cdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy\n",
    "from collections import OrderedDict\n",
    "import zmq\n",
    "import sys\n",
    "import threading\n",
    "import time\n",
    "from random import randint, random\n",
    "import time\n",
    "import zmq\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "numpy.set_printoptions(suppress=False)\n",
    "torch.set_printoptions(sci_mode=False)\n",
    "import tenseal as ts\n",
    "import pickle\n",
    "\n",
    "\n",
    "################################################################\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "from torchsummary import summary\n",
    "import numpy as np\n",
    "import torchvision.utils as vutils\n",
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "\n",
    "ngpu = 1 \n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, ngpu):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.main = nn.Sequential(\n",
    "\n",
    "            nn.Conv2d(1, 32, 4, 2, 1, bias=False),\n",
    "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(32, 64, 4, 2, 1, bias=False),\n",
    "            nn.GroupNorm(64, 64),\n",
    "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(64, 128, 3, 2, 1, bias=False),\n",
    "            nn.GroupNorm(64, 128),\n",
    "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
    "\n",
    "            nn.Conv2d(128, 1, kernel_size=(4, 4), stride=1, bias=False),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "\n",
    "        output = self.main(input)\n",
    "        return output.view(-1, 1).squeeze(1)\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, ngpu):\n",
    "        super(Generator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.main = nn.Sequential(\n",
    "\n",
    "            nn.ConvTranspose2d(100, 128, 4, 1, bias=False),\n",
    "            nn.GroupNorm(32, 128),\n",
    "            nn.ReLU(True),\n",
    "\n",
    "            nn.ConvTranspose2d(128, 64, 3, 2, 1, bias=False),\n",
    "            nn.GroupNorm(32, 64),\n",
    "            nn.ReLU(True),\n",
    "\n",
    "            nn.ConvTranspose2d(64, 32, 4, 2, 1, bias=False),\n",
    "            nn.GroupNorm(32,32),\n",
    "            nn.ReLU(True),\n",
    "\n",
    "            nn.ConvTranspose2d(32, 1, 4, 2, 1, bias=False),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.main(x)\n",
    "\n",
    "\n",
    "\n",
    "G = Generator(ngpu)\n",
    "D = Discriminator(ngpu)\n",
    "\n",
    "# print(G)\n",
    "\n",
    "################################################################\n",
    "\n",
    "D.load_state_dict(torch.load('Dis.ckpt'))\n",
    "discrimiator = D.state_dict()\n",
    "\n",
    "\n",
    "# for k, v in discrimiator.items():\n",
    "#     print(k)\n",
    "    \n",
    "    \n",
    "G.load_state_dict(torch.load('Gen.ckpt'))\n",
    "generator = G.state_dict()\n",
    "\n",
    "# for k, v in generator.items():\n",
    "#     print(k)\n",
    "\n",
    "vals111 = []\n",
    "keys = []\n",
    "vals = []\n",
    "for k, v in discrimiator.items():\n",
    "\n",
    "    keys.append(k)\n",
    "    vals.append(v)\n",
    "    a = v.numpy()\n",
    "    vals111.append(a.shape)\n",
    "    \n",
    "# print(vals111)\n",
    "# vals111 = numpy.array(vals)   \n",
    "# keys = numpy.array(keys)\n",
    "# vals = numpy.array(vals)\n",
    "\n",
    "\n",
    "\n",
    "keys1 = []\n",
    "vals1 = []\n",
    "for k1, v1 in generator.items():\n",
    "\n",
    "    keys1.append(k1)\n",
    "    vals1.append(v1)\n",
    "\n",
    "# keys1 = numpy.array(keys1)\n",
    "# vals1 = numpy.array(vals1)\n",
    "\n",
    "\n",
    "# print(\"Model's state_dict:\")\n",
    "# for param_tensor in model2.state_dict():\n",
    "#     print(torch.numel(model2.state_dict()[param_tensor]))\n",
    "    \n",
    "    \n",
    "\n",
    "###################################################################\n",
    "def elapsed_time_total(start, end):\n",
    "    hours, rem = divmod(end-start, 3600)\n",
    "    minutes, seconds = divmod(rem, 60)\n",
    "    print(\"Total Traning Time: {:0>2}:{:0>2}:{:05.2f}\"\n",
    "                .format(int(hours),int(minutes),seconds))\n",
    "\n",
    "\n",
    "###################################################################\n",
    "\n",
    "\n",
    "import tenseal as ts\n",
    "\n",
    "import base64\n",
    "\n",
    "# Assuming UTF-8 encoding, change to something else if you need to\n",
    "base64.b64encode(\"password\".encode(\"utf-8\"))\n",
    "\n",
    "\n",
    "def write_data(file_name, data):\n",
    "    if type(data) == bytes:\n",
    "        #bytes to base64\n",
    "        data = base64.b64encode(data)\n",
    "         \n",
    "    with open(file_name, 'wb') as f: \n",
    "        f.write(data)\n",
    " \n",
    "def read_data(file_name):\n",
    "    with open(file_name, \"rb\") as f:\n",
    "        data = f.read()\n",
    "    #base64 to bytes\n",
    "    return base64.b64decode(data)\n",
    "\n",
    "\n",
    "###################################################################\n",
    "\n",
    "\n",
    "global data_list\n",
    "global client_num\n",
    "\n",
    "client_num = 0\n",
    "\n",
    "context = zmq.Context()\n",
    "socket = context.socket(zmq.ROUTER)\n",
    "socket.bind(\"tcp://*:5555\")\n",
    "\n",
    "\n",
    "\n",
    "pub_socket = context.socket(zmq.PUB)\n",
    "pub_socket.bind(\"tcp://*:5557\")\n",
    "\n",
    "\n",
    "start_total = time.time()\n",
    "\n",
    "print(\"The server is running now!\")\n",
    "\n",
    "c = 0\n",
    "data_list = []\n",
    "loaded_enc = []\n",
    "loaded_enc_tmp = []\n",
    "cipher1 = []\n",
    "cipher2 = []\n",
    "sum_ = 0\n",
    "\n",
    "data_list_dicriminator = []\n",
    "data_list_generator = []\n",
    "sum_1 = 0\n",
    "sum_2 = 0\n",
    "count = 1\n",
    "while c < 39:\n",
    "    \n",
    "#     print(G)\n",
    "    ident1, msg1 = socket.recv_multipart()\n",
    "    ident2, msg2 = socket.recv_multipart()\n",
    "        \n",
    "    string = b\"New\"\n",
    "    \n",
    "    if string == msg1 and string == msg2: \n",
    "        \n",
    "        client_num = client_num + 1\n",
    "        \n",
    "        message1 = pickle.dumps(vals)\n",
    "        socket.send_multipart([ident1, message1])\n",
    "        \n",
    "        message2 = pickle.dumps(vals1)\n",
    "        socket.send_multipart([ident2, message2])\n",
    "        \n",
    "        print(\"Base model sent to the new client!\")\n",
    "    \n",
    "    else: \n",
    "        \n",
    "        print(\"Training round started\")\n",
    "        \n",
    "        message1 = pickle.loads(msg1)\n",
    "        message2 = pickle.loads(msg2)\n",
    "        \n",
    "        if len(message1) == 8:\n",
    "            \n",
    "            data_list_dicriminator.append(message1)\n",
    "        else:\n",
    "            data_list_generator.append(message1)\n",
    "        \n",
    "        if len(message2) == 8:\n",
    "            data_list_dicriminator.append(message2)\n",
    "        else:\n",
    "            data_list_generator.append(message2)\n",
    "\n",
    "\n",
    "        if len(data_list_dicriminator) == 4 and len(data_list_generator) == 4:\n",
    "            \n",
    "            print(\"Enough data recevied\")\n",
    "            \n",
    "            print(\"len(data_list_dicriminator): \" , len(data_list_dicriminator))\n",
    "            print(\"len(data_list_generator): \" , len(data_list_generator))\n",
    "\n",
    "    \n",
    "            loaded_context = ts.context_from(read_data(\"public.txt\"))\n",
    "        \n",
    "        \n",
    "            for i in range(len(data_list_dicriminator)):\n",
    "                for j in range(len(data_list_dicriminator[0])):  \n",
    "                    loaded_enc_dicriminator = ts.ckks_tensor_from(loaded_context, data_list_dicriminator[i][j])\n",
    "                    seri_loaded_dis = loaded_enc_dicriminator.serialize()\n",
    "                    write_data(f\"loaded_enc_dicriminator{i}{j}.txt\", seri_loaded_dis)\n",
    "            \n",
    "            print(\"dicriminator encrypted loaded\")\n",
    "            \n",
    "            for i in range(len(data_list_generator)):\n",
    "                for j in range(len(data_list_generator[0])):  \n",
    "                    loaded_enc_generator = ts.ckks_tensor_from(loaded_context, data_list_generator[i][j])\n",
    "                    seri_loaded_gen = loaded_enc_generator.serialize()\n",
    "                    write_data(f\"loaded_enc_generator{i}{j}.txt\", seri_loaded_gen)\n",
    "            \n",
    "            print(\"generator encrypted loaded\")\n",
    "            \n",
    "            for i in range(len(data_list_dicriminator[0])): \n",
    "                for j in range(len(data_list_dicriminator)):\n",
    "                    enc_read_dis = read_data(f\"loaded_enc_dicriminator{j}{i}.txt\")\n",
    "                    loaded_enc_dis = ts.ckks_tensor_from(loaded_context, enc_read_dis)                    \n",
    "                    sum_1 += loaded_enc_dis\n",
    "                \n",
    "                sum_1 = sum_1 * 0.25\n",
    "                sum_final1 = sum_1.serialize()\n",
    "                cipher1.append(sum_final1)\n",
    "                sum_1 = 0\n",
    "                sum_final1 = 0\n",
    "            \n",
    "            print(\"Avgg dicriminator encrypted computed\")\n",
    "            \n",
    "            for i in range(len(data_list_generator[0])):  \n",
    "                for j in range(len(data_list_generator)):\n",
    "                    enc_read_gen = read_data(f\"loaded_enc_generator{j}{i}.txt\")\n",
    "                    loaded_enc_gen = ts.ckks_tensor_from(loaded_context, enc_read_gen)\n",
    "                    sum_2 += loaded_enc_gen \n",
    "                sum_2 = sum_2 * 0.25\n",
    "                sum_final2 = sum_2.serialize()\n",
    "                cipher2.append(sum_final2)\n",
    "                sum_2 = 0  \n",
    "                sum_final2 = 0 \n",
    "\n",
    "            print(\"Avgg generator encrypted computed\")\n",
    "            \n",
    "            message1 = pickle.dumps(cipher1)\n",
    "            message2 = pickle.dumps(cipher2)\n",
    "            \n",
    "            pub_socket.send(message1)\n",
    "            pub_socket.send(message2)\n",
    "            \n",
    "            print(\"Sent!\")\n",
    "            print(\"Round: \", count)\n",
    "            count = count + 1\n",
    "            \n",
    "\n",
    "            cipher1 = []\n",
    "            cipher2 = []\n",
    "            sum_1 = 0\n",
    "            sum_2 = 0\n",
    "            sum_final1 = 0\n",
    "            sum_final2 = 0\n",
    "            data_list_dicriminator = []\n",
    "            data_list_generator = []\n",
    "            \n",
    "        c = c + 1\n",
    "        \n",
    "end_total = time.time()\n",
    "elapsed_time_total(start_total, end_total) "
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}