20_StyleGAN_AFHQ.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# StyleGAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is for implementing __A Style-Based Generator Architecture for Generative Adversarial Networks__ from the [paper](https://arxiv.org/abs/1812.04948) with [Tensorflow](https://www.tensorflow.org/).  \n",
    "Datasets are 512x512 [AFHQ](https://github.com/clovaai/stargan-v2).\n",
    "\n",
    "Reference: [tkarras's github](https://github.com/tkarras/progressive_growing_of_gans)(official implementation with tensorflow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IMPORT MODULE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import glob\n",
    "import datetime\n",
    "from sklearn.utils import shuffle\n",
    "import time\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DEFINE PARAMETER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "algorithm = 'StyleGAN_AFHQ'\n",
    "\n",
    "feature_size = 3\n",
    "img_size = [4, 8, 16, 32, 64, 128, 256, 512]\n",
    "gen_fmaps_size = [512, 512, 512, 512, 256, 128, 64, 32]\n",
    "dis_fmaps_size = [512, 512, 512, 512, 512, 256, 128, 64, 32, 16]\n",
    "batch_size = [16, 16, 16, 16, 16, 16, 14, 6]\n",
    "minibatch_size = [16, 16, 16, 16, 16, 16, 14, 6]\n",
    "\n",
    "num_epoch = [27, 54, 54, 54, 54, 54, 54, 54]\n",
    "num_stage = 8\n",
    "start_stage = 0\n",
    "n_latent = 512\n",
    "\n",
    "beta1 = 0\n",
    "beta2 = 0.99\n",
    "ep = 1e-8\n",
    "\n",
    "learning_rate = 0.001\n",
    "\n",
    "train_model = True\n",
    "load_model = False\n",
    "\n",
    "date_time = datetime.datetime.now().strftime(\"%Y%m%d-%H-%M-%S\")\n",
    "save_path = './saved_models/' + date_time + \"_\" + algorithm + \"/save\"\n",
    "load_path = './saved_models/20200511-13-46-12_StyleGAN/save'\n",
    "\n",
    "# WGAN_GP Parateter\n",
    "D_lambda = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IMPORT DATASET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_norm(img):\n",
    "    return (img - (255/2)) / (255/2)\n",
    "\n",
    "def denorm(img):\n",
    "    out = (img + 1) / 2\n",
    "    return np.clip(out, 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_list = []\n",
    "\n",
    "folder_list = os.listdir('./AFHQ/train/')\n",
    "\n",
    "for folder_name in folder_list:\n",
    "    file_list = os.listdir('./AFHQ/train/'+folder_name)\n",
    "    for file_name in file_list:\n",
    "        data_list.append('./AFHQ/train/'+folder_name+'/'+file_name)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_img(nums, stage):\n",
    "    out = []\n",
    "    for num in nums:\n",
    "        out.append(img_norm(cv2.cvtColor(cv2.imread(data_list[num]), cv2.COLOR_BGR2RGB)))      \n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UTILITY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Utility for batch size image\n",
    "def batch_resize(img, img_size):\n",
    "    out = []\n",
    "    for i in range(len(img)):\n",
    "        out.append(cv2.resize(img[i], (img_size, img_size)))\n",
    "    return out\n",
    "\n",
    "def random_latent(batch_size):\n",
    "    return np.clip(np.random.normal(0,1,[batch_size,n_latent]),-1,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NETWORK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class StyleGAN():\n",
    "    def __init__(self, stage):\n",
    "        s = stage\n",
    "        self.alpha = tf.placeholder(shape=[], dtype=tf.float32, name=\"alpha_stage{}\".format(s))\n",
    "        self.latent_c = tf.placeholder(shape=[None, n_latent], dtype=tf.float32, name=\"latent_c_stage{}\".format(s))\n",
    "        self.latent_m = tf.placeholder(shape=[None, n_latent], dtype=tf.float32, name=\"latent_m_stage{}\".format(s))\n",
    "        self.latent_f = tf.placeholder(shape=[None, n_latent], dtype=tf.float32, name=\"latent_f_stage{}\".format(s))\n",
    "        self.img = tf.placeholder(shape=[None, img_size[s], img_size[s], feature_size], dtype=tf.float32, name=\"img_stage{}\".format(s))\n",
    "        \n",
    "        self.fake = self.generator(self.latent_c, self.latent_m, self.latent_f, gen_fmaps_size, s)\n",
    "        \n",
    "        self.d_real = self.discriminator(self.img, dis_fmaps_size, s)\n",
    "        self.d_fake = self.discriminator(self.fake, dis_fmaps_size, s)\n",
    "        \n",
    "        self.var = [var for var in tf.trainable_variables() if \"stage{}\".format(s) in var.name]\n",
    "        self.var_gen = [var for var in self.var if \"generator\" in var.name]\n",
    "        self.var_dis = [var for var in self.var if \"discriminator\" in var.name]\n",
    "        \n",
    "        self.loss_g = self.loss_gen(self.d_fake)\n",
    "        self.loss_d = self.loss_dis(self.img, self.fake, self.d_real, self.d_fake, s)\n",
    "        \n",
    "        opt = tf.train.AdamOptimizer(learning_rate=learning_rate, beta1=beta1, beta2=beta2)\n",
    "        \n",
    "        self.opt_gen = opt.minimize(self.loss_g, var_list=self.var_gen)\n",
    "        self.opt_dis = opt.minimize(self.loss_d, var_list=self.var_dis)\n",
    "        \n",
    "    def get_weight(self, shape):\n",
    "        # not just initializing weights but changing weights dynamically\n",
    "        fan_in = np.prod(shape[:-1])\n",
    "        w_scale = np.sqrt(2 / fan_in) # He initializer constant\n",
    "        w = w_scale * tf.get_variable('weight', shape=shape, initializer=tf.initializers.random_normal())\n",
    "        return w\n",
    "    \n",
    "    def apply_bias_expand(self, x):\n",
    "        b = tf.get_variable('bias', shape=[x.shape[-1]], initializer=tf.initializers.zeros())\n",
    "        return x + tf.reshape(b, [1,1,1,-1])\n",
    "\n",
    "    def apply_bias(self, x):\n",
    "        b = tf.get_variable('bias', shape=[x.shape[-1]], initializer=tf.initializers.zeros())\n",
    "        return x + b\n",
    "    \n",
    "    def dense_expand(self, x, fmaps):\n",
    "        w = self.get_weight([x.shape[1].value, fmaps])\n",
    "        x = tf.matmul(x, w)\n",
    "        return self.apply_bias_expand(x)\n",
    "\n",
    "    def dense(self, x, fmaps):\n",
    "        w = self.get_weight([x.shape[1].value, fmaps])\n",
    "        x = tf.matmul(x, w)\n",
    "        return self.apply_bias(x)\n",
    "    \n",
    "    def conv2d(self, x, fmaps, kernel, padding):\n",
    "        w = self.get_weight([kernel, kernel, x.shape[-1].value, fmaps])\n",
    "        x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding=padding)\n",
    "        return x\n",
    "    \n",
    "    def pixel_norm(self, x, epsilon=1e-8):\n",
    "        # noramlzing weights with axis of feature\n",
    "        return x / tf.sqrt(tf.reduce_mean(tf.square(x), axis=-1, keepdims=True) + epsilon)\n",
    "    \n",
    "    def minibatch(self, x, stage):\n",
    "        group_size = minibatch_size[stage]\n",
    "        s = x.shape\n",
    "        y = tf.reshape(x, [group_size, -1, s[1], s[2], s[3]])\n",
    "        mean = tf.reduce_mean(y, axis=0, keepdims=True) # mean value of minibatch\n",
    "        var = tf.reduce_mean(tf.square(y - mean), axis=0) # variance value of minibatch\n",
    "        std = tf.sqrt(var + 1e-8) # stddev value of minibatch\n",
    "        y = tf.reduce_mean(std, axis=[1,2,3], keepdims=True) \n",
    "        y = tf.tile(y, [group_size, s[1], s[2], 1]) # duplicate value to input size\n",
    "        return tf.concat([x,y], axis=-1)\n",
    "    \n",
    "    def gen_block(self, x, style, fmaps, scale):\n",
    "        s = scale\n",
    "        with tf.variable_scope(\"scale{}\".format(s)):\n",
    "            with tf.variable_scope(\"conv1\"):\n",
    "                x = self.conv2d(x, fmaps[s], 3, \"SAME\")\n",
    "                x = self.apply_noise(x)\n",
    "                x = self.apply_bias_expand(x)\n",
    "                x = tf.nn.leaky_relu(x)\n",
    "                x = self.AdaIN(x, style, fmaps[s])\n",
    "            with tf.variable_scope(\"conv2\"):\n",
    "                x = self.conv2d(x, fmaps[s], 3, \"SAME\")\n",
    "                x = self.apply_noise(x)\n",
    "                x = self.apply_bias_expand(x)\n",
    "                x = tf.nn.leaky_relu(x)\n",
    "                x = self.AdaIN(x, style, fmaps[s])\n",
    "        return x\n",
    "                \n",
    "    def dis_block(self, x, fmaps, scale):\n",
    "        s = scale\n",
    "        with tf.variable_scope(\"scale{}\".format(s)):\n",
    "            with tf.variable_scope(\"conv1\"):\n",
    "                x = self.conv2d(x, fmaps[s+1], 3, \"SAME\")\n",
    "                x = self.apply_bias_expand(x)\n",
    "                x = tf.nn.leaky_relu(x)\n",
    "            with tf.variable_scope(\"conv2\"):\n",
    "                x = self.conv2d(x, fmaps[s], 3, \"SAME\")\n",
    "                x = self.apply_bias_expand(x)\n",
    "                x = tf.nn.leaky_relu(x)\n",
    "        return x\n",
    "    \n",
    "    def upscale(self, x):\n",
    "        size = int(x.shape[1].value * 2)\n",
    "        return tf.image.resize_bilinear(x, [size,size])\n",
    "    \n",
    "    def downscale(self, x):\n",
    "        return tf.nn.avg_pool2d(x, ksize=2, strides=2, padding=\"VALID\")\n",
    "    \n",
    "    def toRGB(self, x, scale, direct):\n",
    "        s = scale\n",
    "        # direct 1 is connected with current scale layer and 2 is connected with previous scale layer\n",
    "        with tf.variable_scope(\"scale{}\".format(s)):\n",
    "            with tf.variable_scope(\"toRGB_{}\".format(direct)):\n",
    "                x = self.conv2d(x, feature_size, 1, \"SAME\")\n",
    "                x = self.apply_bias_expand(x)\n",
    "        return x\n",
    "        \n",
    "    def fromRGB(self, x, fmaps, scale, direct):\n",
    "        s = scale\n",
    "        # direct 1 is connected with current scale layer and 2 is connected with previous scale layer\n",
    "        with tf.variable_scope(\"scale{}\".format(s)):\n",
    "            with tf.variable_scope(\"fromRGB_{}\".format(direct)):\n",
    "                if direct == 1: \n",
    "                    x = self.conv2d(x, fmaps[s+1], 1, \"SAME\")\n",
    "                    x = self.apply_bias_expand(x)\n",
    "                else: \n",
    "                    x = self.conv2d(x, fmaps[s], 1, \"SAME\")\n",
    "                    x = self.apply_bias_expand(x)\n",
    "        return tf.nn.leaky_relu(x)\n",
    "    \n",
    "    def generator(self, z_c, z_m, z_f, fmaps, stage):\n",
    "        with tf.variable_scope(\"generator_stage{}\".format(stage), reuse=tf.AUTO_REUSE):\n",
    "            course_style = self.mapping_network(z_c)\n",
    "            middle_style = self.mapping_network(z_m)\n",
    "            fine_style   = self.mapping_network(z_f)\n",
    "            \n",
    "            with tf.variable_scope(\"scale0\"):\n",
    "                with tf.variable_scope(\"conv1\"):\n",
    "                    x = tf.get_variable('const', shape=[1, 4, 4, fmaps[0]], initializer=tf.initializers.ones())\n",
    "                    x = self.conv2d(x, fmaps[0], 3, \"SAME\")\n",
    "                    x = self.apply_noise(x)\n",
    "                    x = self.apply_bias_expand(x)\n",
    "                    x = tf.nn.leaky_relu(x)\n",
    "                    x = self.AdaIN(x, course_style, fmaps[0])\n",
    "            if stage == 0:\n",
    "                x = self.toRGB(x, 0, 1)\n",
    "\n",
    "            # s th scale layers\n",
    "            for s in range(stage):\n",
    "                if s < 2:\n",
    "                    style = course_style \n",
    "                elif s < 4:\n",
    "                    style = middle_style \n",
    "                else:\n",
    "                    style = fine_style\n",
    "                    \n",
    "                shortcut = x\n",
    "                up1 = self.upscale(x)\n",
    "                x = self.gen_block(up1, style, fmaps, s+1)\n",
    "                if s == stage-1:\n",
    "                    rgb1 = self.toRGB(x, s+1, 1)\n",
    "                    rgb2 = self.toRGB(shortcut, s+1, 2)\n",
    "                    up2 = self.upscale(rgb2)\n",
    "                    x = self.alpha*rgb1 + (1-self.alpha)*up2\n",
    "                    \n",
    "            return tf.nn.tanh(x)\n",
    "    \n",
    "    def discriminator(self, x, fmaps, stage):\n",
    "        with tf.variable_scope(\"discriminator_stage{}\".format(stage), reuse=tf.AUTO_REUSE):\n",
    "            rgb1 = self.fromRGB(x, fmaps, stage, 1)\n",
    "            if stage == 0:\n",
    "                x = rgb1\n",
    "            for s in range(stage):\n",
    "                if s == 0: # the last scale of stage\n",
    "                    down2 = self.downscale(x)\n",
    "                    rgb2 = self.fromRGB(down2, fmaps, stage, 2)\n",
    "                    x = self.dis_block(rgb1, fmaps, stage)\n",
    "                    down1 = self.downscale(x)\n",
    "                    x = self.alpha*down1 + (1-self.alpha)*rgb2\n",
    "                else:\n",
    "                    x = self.dis_block(x, fmaps, stage-s)\n",
    "                    up1 = self.downscale(x)\n",
    "                    x = up1\n",
    "            \n",
    "            mini = self.minibatch(x, stage)\n",
    "    \n",
    "            with tf.variable_scope(\"scale0\"):\n",
    "                with tf.variable_scope(\"conv1\"):\n",
    "                    x = self.conv2d(mini, fmaps[1], 3, \"SAME\")\n",
    "                    x = self.apply_bias_expand(x)\n",
    "                    x = tf.nn.leaky_relu(x)\n",
    "                with tf.variable_scope(\"conv2\"):\n",
    "                    x = self.conv2d(x, fmaps[0], 4, \"VALID\")\n",
    "                    x = self.apply_bias_expand(x)\n",
    "                    x = tf.nn.leaky_relu(x)\n",
    "            \n",
    "                flat = tf.layers.flatten(x)\n",
    "                with tf.variable_scope(\"dense1\"):\n",
    "                    out = self.dense_expand(flat, 1)\n",
    "                    \n",
    "        return out\n",
    "    \n",
    "    # Use WGAN-GP Loss\n",
    "    def loss_gen(self, d_fake):\n",
    "        return -tf.reduce_mean(d_fake)\n",
    "    \n",
    "    def loss_dis(self, real, fake, d_real, d_fake, stage):\n",
    "        loss = tf.reduce_mean(d_fake) - tf.reduce_mean(d_real)\n",
    "        \n",
    "        alpha = tf.random_uniform([tf.shape(real)[0],1,1,1], 0, 1)\n",
    "        mix = alpha*real + (1-alpha)*fake\n",
    "        d_mix = self.discriminator(mix, dis_fmaps_size, stage)\n",
    "        loss_mix = tf.reduce_sum(d_mix)\n",
    "        grad = tf.gradients(loss_mix, [mix])[0]\n",
    "        norm = tf.sqrt(tf.reduce_sum(tf.square(grad), axis=[1,2,3]))\n",
    "        GP = tf.square(norm - 1)\n",
    "        loss += D_lambda*tf.reduce_mean(GP)\n",
    "        \n",
    "        loss += 1e-3*tf.reduce_mean(tf.square(d_real)) # preventing dis_loss from diverging\n",
    "        return loss\n",
    "    \n",
    "    def mapping_network(self, z):\n",
    "        # Normalize, 8 dense \n",
    "        z = self.pixel_norm(z)\n",
    "        with tf.variable_scope('MappingNetwork', reuse=tf.AUTO_REUSE):\n",
    "            for i in range(8):\n",
    "                with tf.variable_scope('Mapping_fc'+str(i), reuse=tf.AUTO_REUSE):\n",
    "                    out = tf.nn.leaky_relu(self.dense(z, n_latent))\n",
    "                    z = out\n",
    "            return out \n",
    "    \n",
    "    # Reference: https://github.com/NVlabs/stylegan\n",
    "    def apply_noise(self, x):\n",
    "        with tf.variable_scope('Noise'):\n",
    "            noise = tf.random_normal([tf.shape(x)[0], x.shape[1], x.shape[2], 1], dtype=x.dtype)\n",
    "            weight = tf.get_variable('weight', shape=[x.shape[-1].value], initializer=tf.initializers.zeros())\n",
    "            return x + noise * tf.reshape(tf.cast(weight, x.dtype), [1, 1, 1, -1])\n",
    "\n",
    "    def AdaIN(self, x, y, n_fmap):\n",
    "        mean, var = tf.nn.moments(x, axes=[1, 2], keep_dims=True)\n",
    "\n",
    "        w = self.MLP(y, n_fmap)\n",
    "\n",
    "        y_s = tf.slice(w, [0, 0], [-1, n_fmap])\n",
    "        y_b = tf.slice(w, [0, n_fmap], [-1, n_fmap])\n",
    "\n",
    "        y_s = tf.expand_dims(tf.expand_dims(y_s, 1), 1)\n",
    "        y_b = tf.expand_dims(tf.expand_dims(y_b, 1), 1)   \n",
    "        \n",
    "        out = y_s * ((x - mean) / var) + y_b \n",
    "        \n",
    "        return out \n",
    "\n",
    "    def MLP(self, x, n_fmap):\n",
    "        with tf.variable_scope('MLP', reuse=tf.AUTO_REUSE):\n",
    "            with tf.variable_scope('fc1', reuse=tf.AUTO_REUSE):\n",
    "                fc1 = tf.nn.leaky_relu(self.dense(x, n_fmap))\n",
    "            with tf.variable_scope('fc2', reuse=tf.AUTO_REUSE):\n",
    "                fc2 = self.dense(fc1, 2*n_fmap)\n",
    "\n",
    "        return fc2\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## INITIALIZATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Initialize variables\n",
    "net = [StyleGAN(s) for s in range(num_stage)] # train stage step by step, 0 to 8\n",
    "init = tf.global_variables_initializer()\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "sess = tf.InteractiveSession(config=config)\n",
    "sess.run(init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOAD MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Load_Saver = [tf.train.Saver(var_list=net[s].var) for s in range(num_stage)]\n",
    "Saver = tf.train.Saver()\n",
    "\n",
    "if load_model == True:\n",
    "    for s in range(1,start_stage):\n",
    "        Load_Saver[s].restore(sess, load_path+str(s))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TRAIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "if train_model:\n",
    "    fixed_latent1 = random_latent(36)\n",
    "    fixed_latent2 = random_latent(36)\n",
    "    len_data = len(data_list)\n",
    "    for s in range(start_stage, num_stage):\n",
    "        alpha = 0\n",
    "        total_batchs = len_data // batch_size[s]\n",
    "        \n",
    "        if s != 0:\n",
    "            for var in net[s].var_gen:\n",
    "                var_gen = var.name.split('/')\n",
    "                for var_old in net[s-1].var_gen:\n",
    "                    var_gen_old = var_old.name.split('/')\n",
    "                    \n",
    "                    var_name_check_list = []\n",
    "                    if len(var_gen) == len(var_gen_old):\n",
    "                        for check_idx in range(1,len(var_gen)):\n",
    "                            var_name_check_list.append(var_gen[check_idx] == var_gen_old[check_idx])\n",
    "                \n",
    "                    if len(var_name_check_list) != 0 and all(var_name_check_list):\n",
    "                        print(var_gen)\n",
    "                        print(var_gen_old)\n",
    "                        sess.run(var.assign(var_old))\n",
    "                    if (\"RGB_2\" in var_gen[2]) and (\"RGB_1\" in var_gen_old[2]) and (var_gen[3] == var_gen_old[3]):\n",
    "                        print(var_gen)\n",
    "                        print(var_gen_old)\n",
    "                        sess.run(var.assign(var_old))\n",
    "            print(\"Generator is assigned!\")\n",
    "\n",
    "            for var in net[s].var_dis:\n",
    "                var_dis = var.name.split('/')\n",
    "                for var_old in net[s-1].var_dis:\n",
    "                    var_dis_old = var_old.name.split('/')\n",
    "                    \n",
    "                    var_name_check_list = []\n",
    "                    if len(var_dis) == len(var_dis_old):\n",
    "                        for check_idx in range(1,len(var_dis)):\n",
    "                            var_name_check_list.append(var_dis[check_idx] == var_dis_old[check_idx])\n",
    "                \n",
    "                    if len(var_name_check_list) != 0 and all(var_name_check_list):\n",
    "                        print(var_dis)\n",
    "                        print(var_dis_old)\n",
    "                        sess.run(var.assign(var_old))\n",
    "                    if (\"RGB_2\" in var_dis[2]) and (\"RGB_1\" in var_dis_old[2]) and (var_dis[3] == var_dis_old[3]):\n",
    "                        print(var_dis)\n",
    "                        print(var_dis_old)\n",
    "                        sess.run(var.assign(var_old))\n",
    "\n",
    "            print(\"Discriminator is assigned!\")\n",
    "            print(\"Assignment is Finished!\")\n",
    "                        \n",
    "        count_batch = 0\n",
    "        for epoch in range(num_epoch[s]):\n",
    "            train_set = shuffle(np.array(range(len_data)))\n",
    "            \n",
    "            for batch in range(total_batchs):\n",
    "                batch_set = train_set[batch * batch_size[s] : (batch + 1) * batch_size[s]]\n",
    "                batch_img = load_img(batch_set, s)\n",
    "#               if stage is 4-64, use batch_resize. Because there is no image dataset for these stages.\n",
    "#               if stage is 128-1024, use original image datasets.\n",
    "                batch_img = batch_resize(batch_img, img_size[s])\n",
    "                batch_latent = random_latent(batch_size[s])\n",
    "\n",
    "                alpha += (1 / (total_batchs*(num_epoch[s]/2)))\n",
    "                if alpha >= 1: alpha = 1\n",
    "                \n",
    "                feed_dict = {net[s].alpha: alpha,\n",
    "                             net[s].latent_c: batch_latent,\n",
    "                             net[s].latent_m: batch_latent,\n",
    "                             net[s].latent_f: batch_latent,\n",
    "                             net[s].img: batch_img}\n",
    "                \n",
    "                _, loss_dis = sess.run([net[s].opt_dis, net[s].loss_d],\n",
    "                                       feed_dict = feed_dict)\n",
    "                \n",
    "                batch_latent = random_latent(batch_size[s])\n",
    "                feed_dict = {net[s].alpha: alpha,\n",
    "                             net[s].latent_c: batch_latent,\n",
    "                             net[s].latent_m: batch_latent,\n",
    "                             net[s].latent_f: batch_latent}\n",
    "                \n",
    "                _, loss_gen = sess.run([net[s].opt_gen, net[s].loss_g],\n",
    "                                              feed_dict = feed_dict)\n",
    "                second = time.time() - start_time\n",
    "                print(\"Stage: {} / Batch: {} / {} / Alpha: {:.3f} / Time: {} d, {} h, {} m, {:.3f} s\".format(\\\n",
    "                    s, count_batch, total_batchs*num_epoch[s], alpha,\\\n",
    "                    int(second // (60*60*24)), int(second // (60*60)) % (24), int(second // 60) % 60, second % 60), end=\"\\r\")\n",
    "                    \n",
    "                count_batch += 1\n",
    "                \n",
    "            fake_test = sess.run(net[s].fake, feed_dict = {net[s].latent_c: fixed_latent1,\n",
    "                                                           net[s].latent_m: fixed_latent1,\n",
    "                                                           net[s].latent_f: fixed_latent1,\n",
    "                                                           net[s].alpha: alpha})\n",
    "            fig, ax = plt.subplots(6,6, figsize=[10,10])\n",
    "            for k in range(36):\n",
    "                ax[k//6,k%6].imshow(denorm(fake_test[k]))\n",
    "                ax[k//6,k%6].axis(\"off\")\n",
    "            plt.tight_layout()\n",
    "            plt.xticks([]), plt.yticks([])\n",
    "            plt.subplots_adjust(left = 0, bottom = 0, right = 1, top = 1, hspace = 0, wspace = 0)\n",
    "            fig.savefig(\"./save_style_progress/\" + str(s) + \"_\" + str(epoch+1).zfill(5) + '.png')\n",
    "            plt.close(fig)\n",
    "            \n",
    "            latent_test1 = random_latent(10)\n",
    "            latent_test2 = random_latent(10)\n",
    "\n",
    "            fake_test_1 = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test1,\n",
    "                                                             net[s].latent_m: latent_test1,\n",
    "                                                             net[s].latent_f: latent_test1,\n",
    "                                                             net[s].alpha: alpha})\n",
    "\n",
    "            fake_test_2 = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test2,\n",
    "                                                             net[s].latent_m: latent_test2,\n",
    "                                                             net[s].latent_f: latent_test2,\n",
    "                                                             net[s].alpha: alpha})\n",
    "\n",
    "            fake_test_c = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test2,\n",
    "                                                             net[s].latent_m: latent_test1,\n",
    "                                                             net[s].latent_f: latent_test1,\n",
    "                                                             net[s].alpha: alpha})\n",
    "\n",
    "            fake_test_m = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test1,\n",
    "                                                             net[s].latent_m: latent_test2,\n",
    "                                                             net[s].latent_f: latent_test1,\n",
    "                                                             net[s].alpha: alpha})\n",
    "\n",
    "            fake_test_f = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test1,\n",
    "                                                             net[s].latent_m: latent_test1,\n",
    "                                                             net[s].latent_f: latent_test2,\n",
    "                                                             net[s].alpha: alpha})\n",
    "\n",
    "            fig, ax = plt.subplots(10,5, figsize=[12,20])\n",
    "\n",
    "            for k in range(10):\n",
    "                ax[k,0].imshow(denorm(fake_test_1[k]))\n",
    "                ax[k,1].imshow(denorm(fake_test_2[k]))\n",
    "                ax[k,2].imshow(denorm(fake_test_c[k]))\n",
    "                ax[k,3].imshow(denorm(fake_test_m[k]))\n",
    "                ax[k,4].imshow(denorm(fake_test_f[k]))\n",
    "\n",
    "                for l in range(5):\n",
    "                    ax[k,l].axis(\"off\")\n",
    "\n",
    "            plt.tight_layout()\n",
    "            plt.xticks([]), plt.yticks([])\n",
    "            plt.subplots_adjust(left = 0, bottom = 0, right = 1, top = 1, hspace = 0, wspace = 0)\n",
    "\n",
    "            fig.savefig(\"./save_style_progress/\" + str(s) + \"_\" + str(epoch+1).zfill(5) + '_style.png')\n",
    "            plt.close(fig)\n",
    "        \n",
    "        Saver.save(sess, save_path+str(s))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TEST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "alpha = 1\n",
    "s = 7\n",
    "latent_test = random_latent(36)\n",
    "fake_test = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test,\n",
    "                                               net[s].latent_m: latent_test,\n",
    "                                               net[s].latent_f: latent_test,\n",
    "                                               net[s].alpha: alpha})\n",
    "fig, ax = plt.subplots(6,6, figsize=[15,15])\n",
    "for k in range(36):\n",
    "    ax[k//6,k%6].imshow(denorm(fake_test[k]))\n",
    "    ax[k//6,k%6].axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.xticks([]), plt.yticks([])\n",
    "plt.subplots_adjust(left = 0, bottom = 0, right = 1, top = 1, hspace = 0, wspace = 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1\n",
    "s = 7\n",
    "\n",
    "latent_test1 = random_latent(10)\n",
    "latent_test2 = random_latent(10)\n",
    "\n",
    "fake_test_1 = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test1,\n",
    "                                                 net[s].latent_m: latent_test1,\n",
    "                                                 net[s].latent_f: latent_test1,\n",
    "                                                 net[s].alpha: alpha})\n",
    "\n",
    "fake_test_2 = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test2,\n",
    "                                                 net[s].latent_m: latent_test2,\n",
    "                                                 net[s].latent_f: latent_test2,\n",
    "                                                 net[s].alpha: alpha})\n",
    "\n",
    "fake_test_c = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test2,\n",
    "                                                 net[s].latent_m: latent_test1,\n",
    "                                                 net[s].latent_f: latent_test1,\n",
    "                                                 net[s].alpha: alpha})\n",
    "\n",
    "fake_test_m = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test1,\n",
    "                                                 net[s].latent_m: latent_test2,\n",
    "                                                 net[s].latent_f: latent_test1,\n",
    "                                                 net[s].alpha: alpha})\n",
    "\n",
    "fake_test_f = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test1,\n",
    "                                                 net[s].latent_m: latent_test1,\n",
    "                                                 net[s].latent_f: latent_test2,\n",
    "                                                 net[s].alpha: alpha})\n",
    "\n",
    "fig, ax = plt.subplots(10,5, figsize=[12,20])\n",
    "\n",
    "for k in range(10):\n",
    "    ax[k,0].imshow(denorm(fake_test_1[k]))\n",
    "    ax[k,1].imshow(denorm(fake_test_2[k]))\n",
    "    ax[k,2].imshow(denorm(fake_test_c[k]))\n",
    "    ax[k,3].imshow(denorm(fake_test_m[k]))\n",
    "    ax[k,4].imshow(denorm(fake_test_f[k]))\n",
    "    \n",
    "    for l in range(5):\n",
    "        ax[k,l].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.xticks([]), plt.yticks([])\n",
    "plt.subplots_adjust(left = 0, bottom = 0, right = 1, top = 1, hspace = 0, wspace = 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1\n",
    "s = 7\n",
    "input_a = random_latent(1)[0]\n",
    "input_b = random_latent(1)[0]\n",
    "latent_test = np.linspace(input_a, input_b, 10)\n",
    "fake_test = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent_test,\n",
    "                                               net[s].latent_m: latent_test,\n",
    "                                               net[s].latent_f: latent_test,\n",
    "                                               net[s].alpha: alpha})\n",
    "\n",
    "fig, ax = plt.subplots(1,10, figsize=[15,15])\n",
    "for k in range(10):\n",
    "    ax[k].imshow(denorm(fake_test[k]))\n",
    "    ax[k].axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.xticks([]), plt.yticks([])\n",
    "plt.subplots_adjust(left = 0, bottom = 0, right = 1, top = 1, hspace = 0, wspace = 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "latent = random_latent(1)\n",
    "fake_test = sess.run(net[s].fake, feed_dict = {net[s].latent_c: latent,\n",
    "                                               net[s].latent_m: latent,\n",
    "                                               net[s].latent_f: latent,\n",
    "                                               net[s].alpha: alpha})\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(denorm(fake_test[0]))\n",
    "plt.axis(\"off\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}