CelebA_gan.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SDDlbqMe5GFh"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ESggJ-VY6eog",
        "outputId": "902d6b00-01f5-4e13-9ff5-edd9105904fd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "O273phIypHZm"
      },
      "outputs": [],
      "source": [
        "\n",
        "!pip --quiet install -U tensorboard-plugin-profile\n",
        "!pip --quiet install wandb\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "E4GA7kz26qBa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6e01e48c-ab1b-4027-b1dc-d9b7c7e274ee"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found GPU at: /device:GPU:0\n"
          ]
        }
      ],
      "source": [
        "\n",
        "from numpy import zeros, ones, asarray\n",
        "from numpy.random import randn, randint\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow.keras.datasets.cifar10 import load_data\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.models import Sequential, save_model, load_model\n",
        "from tensorflow.keras.layers import Dense, Reshape, Flatten, Conv2D, Conv2DTranspose, LeakyReLU, Dropout\n",
        "import tensorflow_datasets as tfds\n",
        "from matplotlib import pyplot as plt\n",
        "from tqdm.notebook import tqdm\n",
        "from tensorflow.keras.models import load_model\n",
        "from numpy.random import randn\n",
        "from time import time\n",
        "from tensorflow.keras import backend as K\n",
        "import numpy as np\n",
        "from itertools import islice\n",
        "\n",
        "import wandb\n",
        "from wandb.keras import WandbCallback\n",
        "\n",
        "import pathlib\n",
        "from google.colab import files\n",
        "\n",
        "\n",
        "import imageio\n",
        "import shutil\n",
        "from os import mkdir, listdir, path, environ\n",
        "import re\n",
        "\n",
        "device_name = tf.test.gpu_device_name()\n",
        "if not device_name:\n",
        "  raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "id": "Rgn2kiMlh2eK"
      },
      "outputs": [],
      "source": [
        "dataset_directory = \"/content/img_align_celeba/\"\n",
        "INPUT_SHAPE = (112,128,3)\n",
        "DATE = \"09-05\"\n",
        "DATASET_SIZE = 202599\n",
        "model_directory = \"/content/drive/MyDrive/models/10epochs\" \n",
        "\n",
        "environ['WANDB_API_KEY '] = \"d136db4e1c22c6659a9bbdee119442d4c564a007\"\n",
        "#!wandb login"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "LP_PATH = \"/content/drive/MyDrive/models/latent_points_dim100\"\n",
        "#LATENT_POINTS = tf.random.normal([5, 100]) \n",
        "#np.save(LP_PATH+\".npy\",LATENT_POINTS.numpy())\n",
        "LATENT_POINTS = np.load(LP_PATH+\".npy\")"
      ],
      "metadata": {
        "id": "C-Uzdx54zr5q"
      },
      "execution_count": 49,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "id": "L0m6vATy1Dhv"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "sweep_config = {\n",
        "    'name':'sweep_1',\n",
        "    'method':'random'\n",
        "}\n",
        "\n",
        "parameters_dic = {\n",
        "    'epochs':{'value':5},\n",
        "    'batch_size':{'values':[64,128]},\n",
        "    'latent_dim':{'values':[50,100,200]},\n",
        "    'gen_struc':{'values':[0,1,2]},\n",
        "    'disc_struc':{'values':[0,1,2]},\n",
        "}\n",
        "sweep_config['parameters'] = parameters_dic\n",
        "#sweep_id = wandb.sweep(sweep_config, project=\"CelebA-GAN-sweep\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "nXpqRiqPFYDJ"
      },
      "outputs": [],
      "source": [
        "!unzip -n -q \"/content/drive/MyDrive/datasets/img_align_celeba.zip\" "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mjEHrC-O9hmY"
      },
      "source": [
        "##Main code"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CGBzRchO5SiI"
      },
      "source": [
        "### Import data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "id": "fbcpEpImaUqb"
      },
      "outputs": [],
      "source": [
        "def make_dataset(batch_size = 64 ):\n",
        "  dataset = tf.keras.preprocessing.image_dataset_from_directory(\n",
        "      dataset_directory, label_mode=None, image_size=(112, 128), batch_size=batch_size,\n",
        "  ) #3min\n",
        "\n",
        "  dataset = dataset.map(lambda x : x / 255.0, num_parallel_calls= tf.data.experimental.AUTOTUNE)\n",
        "  dataset = dataset.map(lambda x : (x*2)-1, num_parallel_calls= tf.data.experimental.AUTOTUNE)\n",
        "  dataset = dataset.prefetch(buffer_size= tf.data.experimental.AUTOTUNE)\n",
        "  return dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Define models"
      ],
      "metadata": {
        "id": "1JGCxP48igNn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "id": "xfQQsJ1x7iBs"
      },
      "outputs": [],
      "source": [
        "\n",
        "def define_discriminator(config,in_shape=INPUT_SHAPE) :\n",
        "  model = Sequential()\n",
        "  if config['disc_struc'] == 0 :\n",
        "    model.add(Conv2D(128, (3,3), strides=(2,2), padding = 'same', input_shape=in_shape))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Conv2D(128, (3,3), strides=(2,2), padding='same'))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Flatten())\n",
        "    model.add(Dropout(0.4))\n",
        "    model.add(Dense(1,activation=\"sigmoid\"))\n",
        "\n",
        "\n",
        "  elif config['disc_struc'] == 1 :\n",
        "    model.add(Conv2D(256, (3,3), strides=(2,2), padding = 'same', input_shape=in_shape))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Conv2D(256, (3,3), strides=(2,2), padding='same'))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Flatten())\n",
        "    model.add(Dropout(0.4))\n",
        "    model.add(Dense(1,activation=\"sigmoid\"))\n",
        "\n",
        "  elif config['disc_struc'] == 2 :\n",
        "    model.add(Conv2D(128, (3,3), strides=(2,2), padding = 'same', input_shape=in_shape))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Conv2D(128, (3,3), strides=(2,2), padding='same'))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Conv2D(128, (3,3), strides=(2,2), padding='same'))\n",
        "    model.add(LeakyReLU(alpha=0.2))\n",
        "\n",
        "    model.add(Flatten())\n",
        "    model.add(Dropout(0.4))\n",
        "    model.add(Dense(1,activation=\"sigmoid\"))\n",
        "\n",
        "  #opt = Adam(learning_rate=0.0002, beta_1=0.5)\n",
        "  opt =  Adam(learning_rate=0.0002, beta_1=0.5)\n",
        "  model.compile(loss=\"binary_crossentropy\", optimizer=opt, metrics=[\"accuracy\"])\n",
        "  return model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "xUl5DZlF-Y0q"
      },
      "outputs": [],
      "source": [
        "def define_generator(config):    #latent_dim is the dimension of the latent vector (e.g., 100)\n",
        "\tmodel = Sequential()\n",
        "\t\n",
        "\tif config['gen_struc'] == 0 :\n",
        "\t\tn_nodes = 64 * 28 * 32  \n",
        "\t\tmodel.add(Dense(n_nodes, input_dim=config['latent_dim'])) #Dense layer so we can work with 1D latent vector\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Reshape((28, 32, 64)))  \n",
        "\t\tmodel.add(Conv2DTranspose(64, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Conv2DTranspose(64, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\telif config['gen_struc'] == 1 : \n",
        "\t\tn_nodes = 512 * 28 * 32  \n",
        "\t\tmodel.add(Dense(n_nodes, input_dim=config['latent_dim'])) #Dense layer so we can work with 1D latent vector\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Reshape((28, 32, 512)))  \n",
        "\t\tmodel.add(Conv2DTranspose(256, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Conv2DTranspose(128, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\telif config['gen_struc'] == 2 : \n",
        "\t\tn_nodes = 128 * 14 * 16  \n",
        "\t\tmodel.add(Dense(n_nodes, input_dim=config['latent_dim'])) #Dense layer so we can work with 1D latent vector\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Reshape((14,16, 128)))  \n",
        "\t\tmodel.add(Conv2DTranspose(128, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Conv2DTranspose(64, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t\tmodel.add(Conv2DTranspose(64, (4,4), strides=(2,2), padding='same'))\n",
        "\t\tmodel.add(LeakyReLU(alpha=0.2))\n",
        "\t# generate\n",
        "\tmodel.add(Conv2D(3, (8,8), activation='tanh', padding='same'))\n",
        "\treturn model  #Model not compiled as it is not directly trained like the discriminator.\n",
        "                    #Generator is trained via GAN combined model. \n",
        "\n",
        "#test_gen = define_generator(config)\n",
        "#print(test_gen.summary())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "id": "NAf_PG6X_n3R"
      },
      "outputs": [],
      "source": [
        "# define the combined generator and discriminator model, for updating the generator\n",
        "#Discriminator is trained separately so here only generator will be trained by keeping\n",
        "#the discriminator constant. \n",
        "def define_gan(generator, discriminator, config):\n",
        "\t\tdiscriminator.trainable = False  #Discriminator is trained separately. So set to not trainable.\n",
        "\t\t# connect generator and discriminator\n",
        "\t\t#opt = Adam(learning_rate=0.0002, beta_1=0.5)\n",
        "\t\topt = Adam(learning_rate=0.0002, beta_1=0.5)\n",
        "\t\tmodel = Sequential()\n",
        "\t\tmodel.add(generator)\n",
        "\t\tmodel.add(discriminator)\n",
        "\t\tmodel.compile(loss='binary_crossentropy', optimizer=opt)\n",
        "\t\treturn model\n",
        "\n",
        "def make_model(config) :\n",
        "  d_model = define_discriminator(config)\n",
        "\t# create the generator\n",
        "  g_model = define_generator(config)\n",
        "\t# create the gan\n",
        "  gan_model = define_gan(g_model, d_model, config)\n",
        "  return gan_model, g_model, d_model\n",
        "\n",
        "#def save(path,gan, generator, discriminator):\n",
        "    #discriminator.trainable = False\n",
        "#    save_model(gan, path+'gan')\n",
        "#    #discriminator.trainable = True\n",
        "#    save_model(generator, path+'generator')\n",
        "#    save_model(discriminator, path+'discriminator')\n",
        "\n",
        "\n",
        "#def load(path):\n",
        "#    discriminator = load_model(path+'discriminator')\n",
        "#    generator = load_model(path+'generator')\n",
        "#    gan = load_model(path+'gan')\n",
        "#    gan.summary()\n",
        "#    discriminator.summary()\n",
        "#    generator.summary()\n",
        "\n",
        "#    return gan, generator, discriminator"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Functions"
      ],
      "metadata": {
        "id": "gvXk5xp2CMll"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def custom_lr_scheduler(epoch, starting_lr) :\n",
        "  if epoch < 10 :\n",
        "    return starting_lr * epoch + starting_lr\n",
        "  if epoch < 30 :\n",
        "    return starting_lr * 10\n",
        "  return starting_lr * 10 * ( 1 - ((epoch-30)/19 ))"
      ],
      "metadata": {
        "id": "vD0Ke8_fzWKW"
      },
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 57,
      "metadata": {
        "id": "WFV5-jYODnht"
      },
      "outputs": [],
      "source": [
        "# pick a batch of random real samples to train the GAN\n",
        "#In fact, we will train the GAN on a half batch of real images and another \n",
        "#half batch of fake images. \n",
        "\n",
        "\n",
        "# generate n_samples number of latent vectors as input for the generator\n",
        "\n",
        "@tf.function\n",
        "def generate_latent_points(latent_dim, n_samples):\n",
        "\t# generate points in the latent space\n",
        "\t#x_input = randn(latent_dim * n_samples)\n",
        "\tx_input = tf.random.normal([n_samples, latent_dim]) \n",
        "\t# reshape into a batch of inputs for the network\n",
        "\t#x_input = x_input.reshape(n_samples, latent_dim)\n",
        "\treturn x_input\n",
        "\n",
        "# use the generator to generate n fake examples, with class labels\n",
        "#Supply the generator, latent_dim and number of samples as input.\n",
        "#Use the above latent point generator to generate latent points. \n",
        "@tf.function\n",
        "def generate_fake_samples(generator, latent_dim, n_samples):\n",
        "\t# generate points in latent space\n",
        "\tx_input = generate_latent_points(latent_dim, n_samples)\n",
        "\t# predict using generator to generate fake samples. \n",
        "\tX = generator(x_input)\n",
        "\t# Class labels will be 0 as these samples are fake. \n",
        "\ty = tf.zeros((n_samples, 1))  #Label=0 indicating they are fake\n",
        "\treturn X, y"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def model_save_directory(path, epoch) : \n",
        "  return path+str(epoch)+\"epochs/\""
      ],
      "metadata": {
        "id": "T-X4m0Qzsdwf"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def atoi(text):\n",
        "    return int(text) if text.isdigit() else text\n",
        "\n",
        "def natural_keys(text):\n",
        "    '''\n",
        "    alist.sort(key=natural_keys) sorts in human order\n",
        "    http://nedbatchelder.com/blog/200712/human_sorting.html\n",
        "    (See Toothy's implementation in the comments)\n",
        "    '''\n",
        "    return [ atoi(c) for c in re.split(r'(\\d+)', text) ]\n",
        "\n",
        "def dl_gif(images_path):\n",
        "  if not path.exists('run_sample'):\n",
        "    mkdir('run_sample')\n",
        "  for i in range(1,6) :\n",
        "    r = re.compile(f'samples_{i}_')\n",
        "    newlist = list(filter(r.match, listdir(images_path)))\n",
        "    if not path.exists(f'run_sample/{i}'):\n",
        "      mkdir(f'run_sample/{i}/')\n",
        "      for image in newlist :  \n",
        "        shutil.move(images_path+image, f'run_sample/{i}/'+image)\n",
        "    filenames = listdir(f'run_sample/{i}/')\n",
        "    filenames.sort(key=natural_keys)\n",
        "    images = []\n",
        "\n",
        "    for filename in filenames:\n",
        "        images.append(imageio.imread(f'run_sample/{i}/'+filename))\n",
        "    imageio.mimsave(f'gif_{i}.gif', images, fps=10)\n",
        "    files.download(f'gif_{i}.gif')"
      ],
      "metadata": {
        "id": "xkeodAavtbpP"
      },
      "execution_count": 59,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_gan(config, model_directory) :\n",
        "  gan_model, g_model, d_model = make_model(config)\n",
        "  latent_points = []\n",
        "  checkpoint = tf.train.Checkpoint(gan_optimizer=gan_model.optimizer,\n",
        "                                    discriminator_optimizer=d_model.optimizer,\n",
        "                                    generator=g_model,\n",
        "                                    discriminator=d_model,\n",
        "                                    gan_model = gan_model,\n",
        "                                    config = config)\n",
        "  checkpoint.restore(model_directory)\n",
        "  return gan_model, g_model, d_model\n",
        "\n",
        "def save_gan(gan_model, g_model, d_model, model_directory, config) :\n",
        "  checkpoint = tf.train.Checkpoint(gan_optimizer=gan_model.optimizer,\n",
        "                                    discriminator_optimizer=d_model.optimizer,\n",
        "                                    generator=g_model,\n",
        "                                    discriminator=d_model,\n",
        "                                    gan_model = gan_model,\n",
        "                                    config = config)\n",
        "  ckpt_manager = tf.train.CheckpointManager(checkpoint, model_directory, max_to_keep=1)\n",
        "  ckpt_manager.save()"
      ],
      "metadata": {
        "id": "MVOt_CJrj6O9"
      },
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Train epoch"
      ],
      "metadata": {
        "id": "oVs9kHNvimgC"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "id": "1eQre8tZDwIf"
      },
      "outputs": [],
      "source": [
        "# train the generator and discriminator\n",
        "#We loop through a number of epochs to train our Discriminator by first selecting\n",
        "#a random batch of images from our true/real dataset.\n",
        "#Then, generating a set of images using the generator. \n",
        "#Feed both set of images into the Discriminator. \n",
        "#Finally, set the loss parameters for both the real and fake images, as well as the combined loss. \n",
        "def train_epoch(g_model, d_model, gan_model, dataset, model_config, training_config, summary_writer):\n",
        "\t# manually enumerate epochs and bacthes. \n",
        "\tlatent_dim = model_config['latent_dim']\n",
        "\tbatch_size = model_config['batch_size']\n",
        "\tnum_epoch = training_config['number_epochs']\n",
        "\n",
        "\tif training_config['sanity_check'] :\n",
        "\n",
        "\t\tds_iterator = [(i,x) for (i,x) in enumerate(tqdm(dataset)) if i<10]\n",
        "\telif training_config['verbose'] == 1 :\n",
        "\t\tds_iterator = enumerate(tqdm(dataset))\n",
        "\telse :\n",
        "\t\tds_iterator = enumerate(dataset)\n",
        "\t\n",
        "\t# enumerate batches over the training set\n",
        "\tfor j, batch in ds_iterator:\n",
        "\t\t#with tf.profiler.experimental.Trace('profiler', step_num=step, _r=1):\n",
        "\t\t\tX_real, y_real = batch, tf.ones([len(batch),1])\n",
        "\t\t\t# update discriminator model weights\n",
        "\t\t\t\t\t\t##train_on_batch allows you to update weights based on a collection \n",
        "\t\t\t\t\t\t#of samples you provide\n",
        "\t\t\t\t\t\t#Let us just capture loss and ignore accuracy value (2nd output below)\n",
        "\t\t\td_loss_real, _ = d_model.train_on_batch(X_real, y_real) \n",
        "\n",
        "\t\t\t\t\t\t# generate 'fake' examples\n",
        "\t\t\tX_fake, y_fake = generate_fake_samples(g_model, latent_dim, batch_size)\n",
        "\t\t\t# update discriminator model weights\n",
        "\t\t\td_loss_fake, _ = d_model.train_on_batch(X_fake, y_fake)\n",
        "\t\t\t\t\t\t\n",
        "\t\t\t\t\t\t#d_loss = 0.5 * np.add(d_loss_real, d_loss_fake) #Average loss if you want to report single..\n",
        "\t\t\t\t\t\t\n",
        "\t\t\t# prepare points in latent space as input for the generator\n",
        "\t\t\tX_gan = generate_latent_points(latent_dim, batch_size)\n",
        "\t\t\t\t\t\t\t\n",
        "\t\t\t\t\t\t# The generator wants the discriminator to label the generated samples\n",
        "\t\t\t\t# as valid (ones)\n",
        "\t\t\t\t#This is where the generator is trying to trick discriminator into believing\n",
        "\t\t\t\t#the generated image is true (hence value of 1 for y)\t\t\t\n",
        "\t\t\ty_gan = ones((batch_size, 1))\n",
        "\t\t\t\t\t\t\n",
        "\t\t\t\t\t\t# Generator is part of combined model where it got directly linked with the discriminator\n",
        "\t\t\t\t# Train the generator with latent_dim as x and 1 as y. \n",
        "\t\t\t\t# Again, 1 as the output as it is adversarial and if generator did a great\n",
        "\t\t\t\t#job of folling the discriminator then the output would be 1 (true)\n",
        "\t\t\t# update the generator via the discriminator's error\n",
        "\t\t\tg_loss1 = gan_model.train_on_batch(X_gan, y_gan)\n",
        "\t \n",
        "\t\t\tX_gan2 = generate_latent_points(latent_dim, batch_size)\n",
        "\t\t\tg_loss2 = gan_model.train_on_batch(X_gan2, y_gan)\n",
        "\t\t\tif summary_writer != None :\n",
        "\t\t\t\twith summary_writer.as_default():\n",
        "\t\t\t\t\ttf.summary.scalar('d_loss_real', d_loss_real, step=j)\n",
        "\t\t\t\t\ttf.summary.scalar('d_loss_fake', d_loss_fake, step=j)\n",
        "\t\t\t\t\ttf.summary.scalar('g_loss', (g_loss1+g_loss2)/2, step=j)\n",
        "\t\t\t\t\t\t# Print losses on this batch\n",
        "\t\t\tif training_config['verbose']  == 2 :\n",
        "\t\t\t\tprint(f'Epoch>{num_epoch+1}, batch={j}, d1={d_loss_real}, d2={d_loss_fake} g={(g_loss1+g_loss2)/2}')\n",
        "\t# save the generator model\n",
        "\n",
        "\tif training_config['verbose']  == 1 :\n",
        "\t\t\t\tprint(f'Epoch>{num_epoch+1}, d1={d_loss_real}, d2={d_loss_fake} g={(g_loss1+g_loss2)/2}')\n",
        "\n",
        "\n",
        "def train(model_config, training_config) :\n",
        "\t#wandb.tensorboard.patch(root_logdir=\"content/log\")\n",
        "\tlatent_points = training_config['latent_points']\n",
        "\tdataset = make_dataset(model_config['batch_size'])\n",
        "\tif training_config['starting_epoch'] == 0 :\n",
        "\t\tgan_model, g_model, d_model = make_model(model_config)\n",
        "\telse :\n",
        "\t\tmodel_dir = model_save_directory(training_config['model_directory'], training_config['starting_epoch'])\n",
        "\t\tmodel_dir = tf.train.latest_checkpoint(model_dir)\n",
        "\t\tgan_model, g_model, d_model = load_gan(model_config,model_dir )\n",
        "\tavg_time = 0\n",
        "\n",
        "\tsummary_writer = tf.summary.create_file_writer('content/log/losses')\n",
        "\n",
        "\tfor epoch in tf.range(training_config['starting_epoch'],training_config['starting_epoch']+training_config['number_epochs']) :\n",
        "\t\tstarting_time = time()\n",
        "\n",
        "\t\tK.set_value(d_model.optimizer.learning_rate, training_config[\"learning_rate_disc\"])\n",
        "\t\tK.set_value(gan_model.optimizer.learning_rate, training_config[\"learning_rate_gan\"])\n",
        "\t\t\n",
        "\t\ttrain_epoch(g_model, d_model, gan_model, dataset, \n",
        "\t\t\t\t\t\t\t\tmodel_config, training_config, \n",
        "\t\t\t\t\t\t\t\tsummary_writer = summary_writer)\n",
        "\t\tdelta_time = time() - starting_time\n",
        "\t\tavg_time += delta_time\n",
        "\t\tX = g_model.predict(latent_points)\n",
        "\t\tX = (X+1)/2\n",
        "\t\tX = (X*255).astype(np.uint8)\n",
        "\t\twandb.log({\"samples_1\": wandb.Image(X[0], caption=f'sample {1}, epoch {epoch}'),\n",
        "\t\t\t\t\t\t\t\"samples_2\": wandb.Image(X[1], caption=f'sample {2}, epoch {epoch}'), \n",
        "\t\t\t\t\t\t\t\"samples_3\": wandb.Image(X[2], caption=f'sample {3}, epoch {epoch}'), \n",
        "\t\t\t\t\t\t\t\"samples_4\": wandb.Image(X[3], caption=f'sample {4}, epoch {epoch}'), \n",
        "\t\t\t\t\t\t\t\"samples_5\": wandb.Image(X[4], caption=f'sample {5}, epoch {epoch}'), \n",
        "\t\t\t\t\t\t\t\t'duration':delta_time,'epoch':epoch+1, \n",
        "\t\t\t\t\t\t\t\t'learning rate discriminator':training_config[\"learning_rate_disc\"],\n",
        "\t\t\t\t\t\t\t\t'learning rate gan':training_config[\"learning_rate_gan\"]\n",
        "\t\t\t\t\t\t\t\t})\n",
        "\twandb.summary['avg_time'] = avg_time/training_config['number_epochs']\n",
        "\tsave_gan(gan_model, g_model, d_model, model_save_directory(training_config['model_directory'], training_config['starting_epoch']+training_config['number_epochs']), model_config)\n",
        "\tif training_config['dl_gif'] :\n",
        "\t\tdl_gif('/content/wandb/latest-run/files/media/images/')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1IJEfB79leR"
      },
      "source": [
        "##Sweep"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {
        "id": "OgwTKnmfC6TX"
      },
      "outputs": [],
      "source": [
        "# Plot generated images \n",
        "def show_plot(examples, n):\n",
        "\tfor i in range(n * n):\n",
        "\t\tplt.subplot(n, n, 1 + i)\n",
        "\t\tplt.axis('off')\n",
        "\t\tplt.imshow(examples[i, :, :, :])\n",
        "\tplt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {
        "id": "5jHmEsw7_XkZ"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "def train_sweep() :\n",
        "  with wandb.init() as run:\n",
        "    config = wandb.config\n",
        "    latent_points = generate_latent_points(config['latent_dim'],5)\n",
        "    dataset = make_dataset(config['batch_size'])\n",
        "    gan_model, g_model, d_model = make_model(config)\n",
        "\n",
        "    avg_time = 0\n",
        "    for epoch in tensorflow.range(config['epochs']) :\n",
        "      starting_time = time()\n",
        "      train_epoch(g_model, d_model, gan_model, dataset, config['latent_dim'],n_batch=config['batch_size'], num_epoch= epoch)\n",
        "      delta_time = time() - starting_time\n",
        "      avg_time += delta_time\n",
        "      X = g_model.predict(latent_points)\n",
        "      X = (X+1)/2\n",
        "      X = (X*255).astype(np.uint8)\n",
        "      wandb.log({\"samples\":[wandb.Image(img, caption=f'sample {id}, epoch {epoch}') for id,img in enumerate(X)],\n",
        "                 'duration':delta_time,'epoch':epoch+1})\n",
        "    wandb.summary['avg_time'] = avg_time/config['epochs']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QN0ZQumHE8iW"
      },
      "source": [
        "## Run sweep"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {
        "id": "rjMFPMBJE-RA"
      },
      "outputs": [],
      "source": [
        "#wandb.agent(\"36gxyv90\", train_sweep, count = 10, project=\"CelebA-GAN-sweep\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Long run"
      ],
      "metadata": {
        "id": "vqoNeAvobWn9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_config = {\n",
        "    'batch_size':64,\n",
        "    'latent_dim':100,\n",
        "    'gen_struc':2,\n",
        "    'disc_struc':2,\n",
        "}\n",
        "training_config = {\n",
        "    'starting_epoch':90,\n",
        "    'number_epochs':10,\n",
        "    'sanity_check':False,\n",
        "    'verbose':1,\n",
        "    'dl_gif':False,\n",
        "    'learning_rate_disc':0.000005,\n",
        "    'learning_rate_gan':0.00004,\n",
        "    'model_directory':\"/content/drive/MyDrive/models/\",\n",
        "    'latent_points':LATENT_POINTS\n",
        "}\n",
        "#wandb.tensorboard.patch(root_logdir=\"content/log\")\n",
        "#wandb.init(config = model_config, project = \"CelebA-GAN-sweep\",sync_tensorboard=True, resume = \"must\", name=\"AXI V1.1\", id=\"2epv65u7\")\n",
        "\n",
        "#train(model_config, training_config)"
      ],
      "metadata": {
        "id": "ue0wUMevbZED"
      },
      "execution_count": 65,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Unit tests"
      ],
      "metadata": {
        "id": "drIln018a6fO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_config = {\n",
        "    'batch_size':64,\n",
        "    'latent_dim':100,\n",
        "    'gen_struc':2,\n",
        "    'disc_struc':2,\n",
        "    'sanity_check':False\n",
        "}\n",
        "training_config = {\n",
        "    'starting_epoch':0,\n",
        "    'number_epochs':1,\n",
        "    'sanity_check':True,\n",
        "    'verbose':1,\n",
        "    'dl_gif':False,\n",
        "    'learning_rate':0.0001,\n",
        "    'model_directory':\"/content/drive/MyDrive/models/\"\n",
        "}\n",
        "#wandb.tensorboard.patch(root_logdir=\"content/log\")\n",
        "wandb.init(config = model_config, project = \"CelebA-GAN-sweep\",sync_tensorboard=True)\n",
        "\n",
        "#train(model_config, training_config)\n",
        "\n",
        "training_config = {\n",
        "    'starting_epoch':1,\n",
        "    'number_epochs':1,\n",
        "    'sanity_check':True,\n",
        "    'verbose':1,\n",
        "    'dl_gif':False,\n",
        "    'learning_rate':0.0001,\n",
        "    'model_directory':\"/content/drive/MyDrive/models/\"\n",
        "}\n",
        "#train(model_config, training_config)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 173,
          "referenced_widgets": [
            "75211c3bd583462a93a701631203810b",
            "8981f6e1569d4939adf30528e3f25fa8",
            "eb1af5136ef84584b94bc67dc7042da7",
            "c83ac45553b8401ca10bf3bb6fb4b98d",
            "3436a02bb63b44dfb2aa419a010641a4",
            "056d39b5f0f440f68ea8eccaab4c2195",
            "3d41f367d2b14002b34870816fbc4a47",
            "720b504df2844d5aa02cb335ddda4c4a"
          ]
        },
        "id": "cnwhavK1a6DI",
        "outputId": "bbac84bd-92f1-4e6d-8adc-b5a685063b35"
      },
      "execution_count": 66,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Finishing last run (ID:yira71v2) before initializing another..."
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "VBox(children=(Label(value='0.001 MB of 0.001 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, max…"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "75211c3bd583462a93a701631203810b"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Synced <strong style=\"color:#cdcd00\">tough-wildflower-1</strong>: <a href=\"https://wandb.ai/keu-3dreconstruction/CelebA-GAN-sweep/runs/yira71v2\" target=\"_blank\">https://wandb.ai/keu-3dreconstruction/CelebA-GAN-sweep/runs/yira71v2</a><br/>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Find logs at: <code>./wandb/run-20220730_125735-yira71v2/logs</code>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Successfully finished last run (ID:yira71v2). Initializing new run:<br/>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Tracking run with wandb version 0.12.21"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Run data is saved locally in <code>/content/wandb/run-20220730_125804-upkbhp58</code>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Syncing run <strong><a href=\"https://wandb.ai/keu-3dreconstruction/CelebA-GAN-sweep/runs/upkbhp58\" target=\"_blank\">swept-yogurt-2</a></strong> to <a href=\"https://wandb.ai/keu-3dreconstruction/CelebA-GAN-sweep\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
            ]
          },
          "metadata": {}
        }
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "SDDlbqMe5GFh",
        "mjEHrC-O9hmY",
        "CGBzRchO5SiI",
        "1JGCxP48igNn",
        "F1IJEfB79leR",
        "QN0ZQumHE8iW"
      ],
      "machine_shape": "hm",
      "name": "CelebA_GAN.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "75211c3bd583462a93a701631203810b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "VBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "VBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "VBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_8981f6e1569d4939adf30528e3f25fa8",
              "IPY_MODEL_eb1af5136ef84584b94bc67dc7042da7"
            ],
            "layout": "IPY_MODEL_c83ac45553b8401ca10bf3bb6fb4b98d"
          }
        },
        "8981f6e1569d4939adf30528e3f25fa8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "LabelModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "LabelModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "LabelView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3436a02bb63b44dfb2aa419a010641a4",
            "placeholder": "​",
            "style": "IPY_MODEL_056d39b5f0f440f68ea8eccaab4c2195",
            "value": "0.009 MB of 0.009 MB uploaded (0.000 MB deduped)\r"
          }
        },
        "eb1af5136ef84584b94bc67dc7042da7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3d41f367d2b14002b34870816fbc4a47",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_720b504df2844d5aa02cb335ddda4c4a",
            "value": 1
          }
        },
        "c83ac45553b8401ca10bf3bb6fb4b98d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3436a02bb63b44dfb2aa419a010641a4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "056d39b5f0f440f68ea8eccaab4c2195": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "3d41f367d2b14002b34870816fbc4a47": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "720b504df2844d5aa02cb335ddda4c4a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}