FedGAN-HE-DP-MNIST-Client.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae1f93a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy\n",
    "import pickle\n",
    "from collections import OrderedDict\n",
    "import json\n",
    "import zmq\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import numpy\n",
    "import sys\n",
    "from uuid import uuid4\n",
    "import uuid\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",
    "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",
    "def is_pickle_stream(stream):\n",
    "    try:\n",
    "        pickle.loads(stream)\n",
    "        return True\n",
    "    except:\n",
    "        return False\n",
    "\n",
    "context2 = zmq.Context()\n",
    "print(\"Connecting to key server ...\")\n",
    "socket = context2.socket(zmq.DEALER)\n",
    "socket.connect(\"tcp://localhost:5555\")\n",
    "identity = str(id)\n",
    "socket.identity = identity.encode(\"ascii\")\n",
    "    \n",
    "\n",
    "\n",
    "context1 = zmq.Context()\n",
    "print(\"Connecting to hello world server…\")\n",
    "socket1 = context1.socket(zmq.DEALER)\n",
    "socket1.connect(\"tcp://localhost:5556\")\n",
    "identity = str(id)\n",
    "socket1.identity = identity.encode(\"ascii\")\n",
    "    \n",
    "    \n",
    "sub_socket = context2.socket(zmq.SUB)\n",
    "sub_socket.connect(\"tcp://localhost:5557\")\n",
    "sub_socket.setsockopt_string(zmq.SUBSCRIBE, '')\n",
    "\n",
    "\n",
    "\n",
    "def train(discriminator_arrived, generator_arrived):\n",
    "\n",
    "    import random\n",
    "    import numpy as np\n",
    "    import torch.backends.cudnn as cudnn\n",
    "    import torch.nn as nn\n",
    "    import torch.optim as optim\n",
    "    import torch.utils.data\n",
    "    from torchvision import datasets, transforms\n",
    "    import torchvision.transforms as transforms\n",
    "    import torchvision.utils as vutils\n",
    "    from opacus import PrivacyEngine\n",
    "    from tqdm import tqdm\n",
    "    import matplotlib.pyplot as plt\n",
    "    import time\n",
    "    \n",
    "    nc = 1\n",
    "    \n",
    "    workers = 2\n",
    "    batch_size = 128\n",
    "    imageSize = 28\n",
    "\n",
    "\n",
    "    epochs = 20\n",
    "    lr = 0.0002\n",
    "    beta1 = 0.5\n",
    "    ngpu = 1\n",
    "\n",
    "    target_digit = 9\n",
    "\n",
    "    nc = 1\n",
    "    \n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "    disable_dp = False\n",
    "\n",
    "    secure_rng = False\n",
    "\n",
    "    r = 1\n",
    "\n",
    "    n_runs = 1\n",
    "\n",
    "    sigma = 0.5\n",
    "\n",
    "    max_per_sample_grad_norm = 1.0\n",
    "\n",
    "    delta = 1e-4\n",
    "\n",
    "    nz = 100\n",
    "\n",
    "\n",
    "    def elapsed_time(start, end):\n",
    "        hours, rem = divmod(end-start, 3600)\n",
    "        minutes, seconds = divmod(rem, 60)\n",
    "        print(\"Elapsed Time: {:0>2}:{:0>2}:{:05.2f}\"\n",
    "                    .format(int(hours),int(minutes),seconds))\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",
    "    manualSeed = random.randint(1, 50000)\n",
    "\n",
    "    print(\"Random Seed: \", manualSeed)\n",
    "    random.seed(manualSeed)\n",
    "    torch.manual_seed(manualSeed)\n",
    "\n",
    "    \n",
    "    try:\n",
    "        dataset = datasets.MNIST(\n",
    "            root='/PP-FEDGAN/Data',\n",
    "            download=True,\n",
    "            transform=transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize(imageSize),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize((0.5,), (0.5,)),\n",
    "                ]\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    except ValueError:\n",
    "        print(\"Cannot load dataset\")\n",
    "\n",
    "\n",
    "    dataset_range = list(range(0, 5000))\n",
    "\n",
    "    trainset_range = torch.utils.data.Subset(dataset, dataset_range)\n",
    "\n",
    "    dataloader = torch.utils.data.DataLoader(\n",
    "        trainset_range,\n",
    "        num_workers=int(workers),\n",
    "        batch_size=batch_size,\n",
    "    )\n",
    "\n",
    "    \n",
    "    \n",
    "    \n",
    "    def weights_init(m):\n",
    "        classname = m.__class__.__name__\n",
    "        if classname.find(\"Conv\") != -1:\n",
    "            m.weight.data.normal_(0.0, 0.02)\n",
    "        elif classname.find(\"BatchNorm\") != -1:\n",
    "            m.weight.data.normal_(1.0, 0.02)\n",
    "            m.bias.data.fill_(0)\n",
    "\n",
    "    def elapsed_time(start, end):\n",
    "        hours, rem = divmod(end-start, 3600)\n",
    "        minutes, seconds = divmod(rem, 60)\n",
    "        print(\"Elapsed Time: {:0>2}:{:0>2}:{:05.2f}\"\n",
    "                    .format(int(hours),int(minutes),seconds))\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",
    "    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",
    "    G = G.to(device)\n",
    "    \n",
    "    D = Discriminator(ngpu)\n",
    "    D = D.to(device)\n",
    "\n",
    "    fixed_noise = torch.randn(batch_size, 100, 1, 1, device=device)\n",
    "\n",
    "    G.apply(weights_init)\n",
    "    D.apply(weights_init)\n",
    "    \n",
    "    REAL_LABEL = 1.0\n",
    "    FAKE_LABEL = 0.0\n",
    "   ####################################################### \n",
    "    \n",
    "    discriminator_old = D.state_dict()\n",
    "    shape_d = []\n",
    "    for name, param in discriminator_old.items():\n",
    "        \n",
    "        shape_d.append(param.cpu().numpy())\n",
    "        \n",
    "    \n",
    "#     discriminator_arrived, generator_arrived\n",
    "\n",
    "    valss_d = []\n",
    "    de_ser_d = []\n",
    "    \n",
    "#     if isinstance(discriminator_arrived[0], (bytes, bytearray)):\n",
    "#         print(\"byte\")    \n",
    "#         for i in range(len(discriminator_arrived)):\n",
    "#             loaded_enc_dis_1 = ts.ckks_tensor_from(context, discriminator_arrived[i])\n",
    "#             final_dd = (loaded_enc_dis_1.decrypt().tolist())\n",
    "#             print(len(final_dd))\n",
    "    \n",
    "    if isinstance(discriminator_arrived[0], (bytes, bytearray)):\n",
    "        print(\"byte\")\n",
    "        for i in range(len(discriminator_arrived)):\n",
    "            loaded_enc_dis = ts.ckks_tensor_from(context, discriminator_arrived[i])\n",
    "            final_d = (loaded_enc_dis.decrypt().tolist())\n",
    "            final__d = (numpy.reshape(final_d, shape_d[i].shape ))\n",
    "            valss_d.append(torch.from_numpy(final__d))\n",
    "#             \n",
    "    else:\n",
    "        print(\"Not byte\")\n",
    "        valss_d = discriminator_arrived\n",
    "    \n",
    "    generator_old = G.state_dict()\n",
    "    shape_g = []\n",
    "    \n",
    "    for name, param in generator_old.items():\n",
    "        \n",
    "        shape_g.append(param.cpu().numpy())\n",
    "    \n",
    "    valss_g = []\n",
    "    de_ser_g = []\n",
    "    \n",
    "    if isinstance(generator_arrived[0], (bytes, bytearray)):\n",
    "        print(\"byte\")\n",
    "        for i in range(len(generator_arrived)):\n",
    "            loaded_enc_gen = ts.ckks_tensor_from(context, generator_arrived[i])\n",
    "            final_g = (loaded_enc_gen.decrypt().tolist())\n",
    "#             print(len(final))\n",
    "            final__g = (numpy.reshape(final_g, shape_g[i].shape ))\n",
    "            valss_g.append(torch.from_numpy(final__g))\n",
    "#         print(valss)\n",
    "#             \n",
    "    else:\n",
    "        print(\"Not byte\")\n",
    "        valss_g = generator_arrived\n",
    "    \n",
    "    \n",
    "\n",
    "    #######################################################\n",
    "    i = 0\n",
    "\n",
    "    for name, param in discriminator_old.items():\n",
    "        # Don't update if this is not a weight.\n",
    "    #     if not \"weight\" in name:\n",
    "    #         continue\n",
    "\n",
    "        # Transform the parameter as required.\n",
    "        transformed_param_d = valss_d[i]\n",
    "\n",
    "        # Update the parameter.\n",
    "        param.copy_(transformed_param_d) \n",
    "        i = i + 1\n",
    "        \n",
    "        \n",
    "    j = 0\n",
    "\n",
    "    for name, param in generator_old.items():\n",
    "        # Don't update if this is not a weight.\n",
    "    #     if not \"weight\" in name:\n",
    "    #         continue\n",
    "\n",
    "        # Transform the parameter as required.\n",
    "        transformed_param_g = valss_g[j]\n",
    "\n",
    "        # Update the parameter.\n",
    "        param.copy_(transformed_param_g) \n",
    "        j = j + 1 \n",
    "        \n",
    "    #######################################################\n",
    "\n",
    "    criterion = nn.BCELoss()\n",
    "\n",
    "    optim_D = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))\n",
    "\n",
    "    if not disable_dp:\n",
    "\n",
    "        privacy_engine = PrivacyEngine(D)\n",
    "\n",
    "        D, optim_D, dataloader = privacy_engine.make_private(\n",
    "\n",
    "            module=D,\n",
    "            optimizer=optim_D,\n",
    "            data_loader=dataloader,\n",
    "            noise_multiplier=sigma,\n",
    "            max_grad_norm=max_per_sample_grad_norm,\n",
    "        )\n",
    "\n",
    "    optim_G = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))\n",
    "\n",
    "\n",
    "    \n",
    "    start_total = time.time()\n",
    "    iters = 0\n",
    "    img_list = []\n",
    "    fixed_noise1 = torch.randn(16, 100, 1, 1).cuda()\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        \n",
    "        start = time.time()\n",
    "        \n",
    "        data_bar = tqdm(dataloader)\n",
    "        \n",
    "        for i, data in enumerate(data_bar, 0):\n",
    "            ############################\n",
    "            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "            ###########################\n",
    "\n",
    "            optim_D.zero_grad(set_to_none=True)\n",
    "\n",
    "            real_data = data[0].to(device)\n",
    "            batch_size = real_data.size(0)\n",
    "\n",
    "            # train with real\n",
    "            label_true = torch.full((batch_size,), REAL_LABEL, device=device)\n",
    "            output = D(real_data)\n",
    "            errD_real = criterion(output, label_true)\n",
    "            D_x = output.mean().item()\n",
    "\n",
    "            # train with fake\n",
    "            noise = torch.randn(batch_size, nz, 1, 1, device=device)\n",
    "            fake = G(noise)\n",
    "            label_fake = torch.full((batch_size,), FAKE_LABEL, device=device)\n",
    "            output = D(fake.detach())\n",
    "            errD_fake = criterion(output, label_fake)\n",
    "\n",
    "            # below, you actually have two backward passes happening under the hood\n",
    "            # which opacus happens to treat as a recursive network\n",
    "            # and therefore doesn't add extra noise for the fake samples\n",
    "            # noise for fake samples would be unnecesary to preserve privacy\n",
    "\n",
    "            errD = errD_real + errD_fake\n",
    "            errD.backward()\n",
    "            optim_D.step()\n",
    "            optim_D.zero_grad(set_to_none=True)\n",
    "\n",
    "            D_G_z1 = output.mean().item()\n",
    "\n",
    "            ############################\n",
    "            # (2) Update G network: maximize log(D(G(z)))\n",
    "            ###########################\n",
    "            optim_G.zero_grad()\n",
    "\n",
    "            label_g = torch.full((batch_size,), REAL_LABEL, device=device)\n",
    "            output_g = D(fake)\n",
    "            errG = criterion(output_g, label_g)\n",
    "            errG.backward()\n",
    "            D_G_z2 = output_g.mean().item()\n",
    "            optim_G.step()\n",
    "\n",
    "            if not disable_dp:\n",
    "                epsilon, best_alpha = privacy_engine.accountant.get_privacy_spent(\n",
    "                    delta=delta\n",
    "                )\n",
    "                data_bar.set_description(\n",
    "                    f\"epoch: {epoch}, Loss_D: {errD.item()} \"\n",
    "                    f\"Loss_G: {errG.item()} D(x): {D_x} \"\n",
    "                    f\"D(G(z)): {D_G_z1}/{D_G_z2}\"\n",
    "                    \"(ε = %.2f, δ = %.2f) for α = %.2f\" % (epsilon, delta, best_alpha)\n",
    "                )\n",
    "            else:\n",
    "                data_bar.set_description(\n",
    "                    f\"epoch: {epoch}, Loss_D: {errD.item()} \"\n",
    "                    f\"Loss_G: {errG.item()} D(x): {D_x} \"\n",
    "                    f\"D(G(z)): {D_G_z1}/{D_G_z2}\"\n",
    "                )\n",
    "                \n",
    "        from imutils import build_montages\n",
    "        import os\n",
    "        import cv2\n",
    "\n",
    "#         benchmarkNoise = torch.randn(256, 100, 1, 1, device=device)\n",
    "\n",
    "            # set the generator in evaluation phase, make predictions on\n",
    "            # the benchmark noise, scale it back to the range [0, 255],\n",
    "            # and generate the montage\n",
    "        \n",
    "        from random import randrange\n",
    "        c= randrange(1000)\n",
    "        \n",
    "#         if (iters % 500 == 0):\n",
    "            \n",
    "#             G.eval()\n",
    "#             images = G(benchmarkNoise)\n",
    "#             images = images.detach().cpu().numpy().transpose((0, 2, 3, 1))\n",
    "#             images = ((images * 127.5) + 127.5).astype(\"uint8\")\n",
    "#             images = np.repeat(images, 3, axis=-1)\n",
    "#             vis = build_montages(images, (64, 64), (8, 4))[0]\n",
    "#             p = os.path.join('/content', \"epoch_{}_{}.png\".format(str(epoch + 1).zfill(4),c))\n",
    "#             cv2.imwrite(p, vis)\n",
    "            \n",
    "        with torch.no_grad():\n",
    "            fake = G(fixed_noise1).to(device).cpu().detach()\n",
    "        img_list.append(vutils.make_grid(fake,nrow = 8, normalize=True))\n",
    "\n",
    "            \n",
    "        iters += 1\n",
    "        \n",
    "        end = time.time()\n",
    "        elapsed_time(start, end)\n",
    "        \n",
    "        fig = plt.figure(figsize=(5, 5))\n",
    "        rand_noise = torch.rand((64, 100, 1, 1))\n",
    "        out = vutils.make_grid(G(rand_noise.to(device)).cpu().detach(), padding=5, normalize=True)\n",
    "        plt.imshow(np.transpose(out.numpy(), (1, 2, 0)), cmap=\"gray\")\n",
    "        plt.show()\n",
    "        \n",
    "    fig = plt.figure(figsize=(10, 10))\n",
    "    plt.subplot(1,2,2)\n",
    "    plt.axis(\"off\")\n",
    "#     plt.title(sigma)\n",
    "#     plt.title('Scatter plot pythonspot.com', y=-0.01)\n",
    "    plt.imshow(np.transpose(img_list[-1].numpy(),(1,2,0)))\n",
    "    plt.show()\n",
    "    \n",
    "    print(iters)\n",
    "    \n",
    "#     from IPython.display import Image\n",
    "#     from torchvision.utils import save_image\n",
    "#     import os\n",
    "\n",
    "\n",
    "#     sample_vectors = torch.rand(100, 100, 1, 1).to(device)\n",
    "#     sample_dir = \"/MNIST_DP_TEST/Data/1.4/\"\n",
    "\n",
    "#     def denorm(x):\n",
    "#         out = (x + 1) / 2\n",
    "#         return out.clamp(0, 1)\n",
    "\n",
    "\n",
    "#     def save_fake_images(index, fixed_noise11):\n",
    "#         fake_images = G(fixed_noise11)\n",
    "#         fake_images = fake_images.reshape(fake_images.size(0), 1, 28, 28)\n",
    "#         fake_fname = 'fake_images-{0:0=4d}.png'.format(index)\n",
    "#         print('Saving', fake_fname)\n",
    "#         save_image(denorm(fake_images), os.path.join(sample_dir, fake_fname), nrow=10)\n",
    "\n",
    "#     for i in range (600):\n",
    "#         torch.manual_seed(i)\n",
    "#         fixed_noise11 = torch.randn(100, 100, 1, 1).cuda()\n",
    "#     # Before training\n",
    "#         save_fake_images(i,fixed_noise11)\n",
    "#     # Image(os.path.join(sample_dir, 'fake_images-0000.png'))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    #######################################################\n",
    "\n",
    "    model_last_d = D.state_dict()\n",
    "    \n",
    "    keys_d = []\n",
    "    vals_d = []\n",
    "    \n",
    "    for k, v in model_last_d.items():\n",
    "        \n",
    "        keys_d.append(k)\n",
    "        vals_d.append(v.cpu())\n",
    "        \n",
    "#     keys_d = numpy.array(keys_d)\n",
    "#     vals_d = numpy.array(vals_d)\n",
    "    \n",
    "    cipher_d = []\n",
    "    for i in range(len(vals_d)):\n",
    "        print(vals_d[i].shape)\n",
    "        encrypted_tensor_d = ts.ckks_tensor(context, vals_d[i], batch = True)\n",
    "        print(encrypted_tensor_d)\n",
    "        ser_tensor_d = encrypted_tensor_d.serialize()\n",
    "        print(len(ser_tensor_d))\n",
    "        cipher_d.append(ser_tensor_d)\n",
    "        print(len(cipher_d))\n",
    "\n",
    "\n",
    "\n",
    "    model_last_g = G.state_dict()\n",
    "    \n",
    "    keys_g = []\n",
    "    vals_g = []\n",
    "    \n",
    "    for h, z in model_last_g.items():\n",
    "\n",
    "        keys_g.append(h)\n",
    "        vals_g.append(z.cpu())\n",
    "\n",
    "#     keys_g = numpy.array(keys_g)\n",
    "#     vals_g = numpy.array(vals_g)\n",
    "    \n",
    "    cipher_g = []\n",
    "    for i in range(len(vals_g)):\n",
    "        encrypted_tensor_g = ts.ckks_tensor(context, vals_g[i], batch = True)\n",
    "        ser_tensor_g = encrypted_tensor_g.serialize()\n",
    "        cipher_g.append(ser_tensor_g)\n",
    "\n",
    "    #######################################################\n",
    "    \n",
    "    end_total = time.time()\n",
    "    elapsed_time_total(start_total, end_total)\n",
    "\n",
    "    return cipher_d , cipher_g\n",
    "\n",
    "\n",
    "context = ts.context_from(read_data(\"secret.txt\"))\n",
    "\n",
    "for request in range(10):\n",
    "    \n",
    "    print(\"request in loop: \" , request)\n",
    "    \n",
    "    if request == 0 :\n",
    "        \n",
    "        print(\"request in if\" , request)\n",
    "        message = b\"New\"\n",
    "        socket.send(message)\n",
    "        socket.send(message)\n",
    "\n",
    "        discriminator = socket.recv()\n",
    "        generator = socket.recv()\n",
    "        discriminator = pickle.loads(discriminator) \n",
    "        generator = pickle.loads(generator)\n",
    "        \n",
    "    else: \n",
    "        print(\"request in else\" , request)\n",
    "        discriminator = sub_socket.recv()\n",
    "        generator = sub_socket.recv()\n",
    "        print(f\"New update number {request} received from the server \")\n",
    "    \n",
    "#     ident,  message = socket.recv_multipart()\n",
    "    if is_pickle_stream(discriminator) and is_pickle_stream(generator):\n",
    "        discriminator = pickle.loads(discriminator) \n",
    "        generator = pickle.loads(generator) \n",
    "        print(\"if train started\")\n",
    "        cipher_d, cipher_g = train(discriminator, generator)\n",
    "        print(\"len(cipher_d):   \", len(cipher_d))\n",
    "        print(\"len(cipher_g):   \", len(cipher_g))\n",
    "    else: \n",
    "        print(\"else train started\")\n",
    "        cipher_d, cipher_g = train(discriminator, generator)\n",
    "        print(\"len(cipher_d):   \", len(cipher_d))\n",
    "        print(\"len(cipher_g):   \", len(cipher_g))\n",
    "    print(\"train finished\")\n",
    "    print(f\"Model number {request} locally trained\")\n",
    "\n",
    "    cipher_d_ = pickle.dumps(cipher_d) \n",
    "    cipher_g_ = pickle.dumps(cipher_g) \n",
    "    socket.send(cipher_d_)\n",
    "    socket.send(cipher_g_)"
   ]
  }
 ],
 "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
}