VAE+GAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Автоэнкодеры в Keras\n",
    "\n",
    "\n",
    "# Часть 6: VAE + GAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Содержание\n",
    "\n",
    "* Часть 1: Введение\n",
    "* Часть 2: *Manifold learning* и скрытые (*latent*) переменные\n",
    "* Часть 3: Вариационные автоэнкодеры (*VAE*)\n",
    "* Часть 4: *Conditional VAE*\n",
    "* Часть 5: *GAN* (Generative Adversarial Networks) и tensorflow\n",
    "* **Часть 6: *VAE* + *GAN* **\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В позапрошлой части мы создали ***CVAE*** автоэнкодер, декодер которого умеет генерировать цифру заданного лейбла, мы также попробовали создавать картинки цифр других лейблов в стиле заданной картинки. Получилось довольно хорошо, однако цифры генерировались смазанными.  \n",
    "В прошлой части мы изучили как работают ***GAN'ы***, получив довольно четкие изображения цифр, однако пропала возможность кодирования и переноса стиля.\n",
    "\n",
    "В этой части попробуем взять лучшее от обоих подходов путем совмещения *вариационных автоэнкодеров* (***VAE***) и *генеративных состязающихся сетей* (***GAN***). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Подход который будет описан далее основан на статье ***[Autoencoding beyond pixels using a learned similarity metric, Larsen et al, 2016]***."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Разберемся более подробно почему восстановленные изображения получаются смазанные.  \n",
    "В части про *VAE* рассматривался процесс генерации изображений $X$ из скрытых (*latent*) переменных $Z$.   \n",
    "Так как размерность скрытых переменных $Z$ значительно ниже, чем размерность объектов $X$ (в части про *VAE* эти размерности были 2 и 784), то одному и тому же $Z$ может соответствовать многомерное распределение $X$, то есть $P(X|Z)$. Это распределение можно представить как:\n",
    "\n",
    "$$\n",
    "P(X|Z) = f(Z) + \\epsilon,\n",
    "$$\n",
    "\n",
    "где $f(Z)$ некоторый средний наиболее вероятный объект при заданном $Z$, а $\\epsilon$ - шум какой-то сложной природы.\n",
    "\n",
    "Когда мы обучаем автоэнкодеры, мы сравниваем вход из выборки $X_s$ и выход автоэнкодера $\\tilde X_s$ с помощью некоторого функционала ошибки $L$,\n",
    "\n",
    "$$\n",
    "L(X_s, \\tilde X_s), \\\\\n",
    "\\tilde X_s = f_d(Z; \\theta_d), \\\\\n",
    "Z \\sim Q(Z|X_s; \\theta_e),\n",
    "$$\n",
    "\n",
    "где $Q,\\ f_d$ - энкодер и декодер.\n",
    "\n",
    "Задавая $L$ - мы определяем шум $\\epsilon_L$, которым приближаем настоящий шум $\\epsilon$.  \n",
    "Минимизируя $L$ мы учим автоэнкодер подстраиваться под шум $\\epsilon_L$, убирая его, то есть находить среднее значение в заданной метрике (во второй части это показывалось наглядно на простом искусственном примере)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Если шум $\\epsilon_L$, который мы определеям функционалом $L$ не соответствует рельному шуму $\\epsilon$, то $f_d(Z; \\theta_2)$ окажется сильно смещеным от реального $f(Z)$ (пример: если в регрессии реальный шум лаплассовский, а минимизируется разность квадратов, то предсказанное значение будет смещено в сторону выбросов). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Возвращаясь к картинкам: посмотрим как связана попиксельная метрика которой определен лосс в предыдущих частях, и метрика используемая человеком. Пример и иллюстрация из ***[2]***  \n",
    "\n",
    "![](figs6/metric.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "На картинке выше: (а) - оригинальное изображение цифры, (b) - получена из (а) отрезанием куска, (с) - цифра (а) сдвинутая на пол пикселя вправо.  \n",
    "С точки зрения попиксельной метрики: (а) намного ближе к (b), чем к (с);  \n",
    "хотя с точки зрения человеческого восприятия (b) - даже не цифра, а вот разница между (а) и (b) практически незаметна."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Автоэнкодеры с попиксельной метрикой таким образом размазывали изображение, отражая тот факт, что \n",
    "в рамках близких $Z$:  \n",
    "- положение цифр слегка гуляет по картинке,  \n",
    "- нарисованы цифры слегка поразному (хотя попиксельно может быть значительно далеко)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "По метрике же человеческого восприятия, тот факт, что цифра размылась уже заставляет ее быть сильно не похожей на оригинал. Таким образом если мы будем знать метрику человека или близкую к ней и оптимизировать в ней, то цифры не будут размываться, а важность того, чтобы цифра была полноценной, не как с картинки (b), резко возрастет."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Можно пытаться вручную придумывать метрику, которая будет ближе к человеческой. Но используя подход ***GAN*** можно обучить нейронную сеть самой искать хорошую метрику."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Про *GAN'ы* написано в прошлой части."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Соединяя *VAE* и *GAN*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Генератор *GAN* выполняет функцию аналогичную декодеру в *VAE*: оба сэмплят из априорного распределения $P(Z)$ и переводят его в $P_g(X)$. Однако роли у них разные: декодер восстанавливает объект закодированный энкодером, при обучении опираясь на некоторую метрику сравнения; генератор же генерирует случайный объект, который ни с чем не сравнивается, лишь бы дискриминатор не мог отличить какому из распределений $P$ или $P_g$ он принадлежит."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Идея: добавить в *VAE* третью сеть - дискриминатор, и подавать ей на вход и восстановленный объект и оригинал, а дискриминатор обучать определять какой из них какой.\n",
    "![](figs6/VAE_GAN_scheme.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Иллюстрация из ***[1]***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Разумеется использовать ту же самую метрику сравнения из *VAE* мы уже не можем, потому что обучаясь в ней декодер генерирует изображения легко отличимые от оригинала. Не использовать метрику вообще - тоже, так как нам бы хотелось, чтобы воссозданный $\\hat X$ был похож на оригинал, а не просто какой-то случайный из $P(X)$, как в чистом *GAN*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Задумаемся, однако, вот о чем: дискриминатор, учась отличать реальный объект от сгенерированного, будет вычлинять какие-то характерные черты одних и других. Эти черты объекта будут закодированы в слоях дискриминатора и на основе их комбинации он уже будет выдавать вероятность объекта быть реальным. Например, если изображение размыто, то какой-то нейрон в дискриминаторе будет активироваться сильнее, чем если оно четкое. При этом чем глубже слой, тем более абстрактные характеристики входного объекта в нем закодированы.\n",
    "\n",
    "Так как каждый слой дискриминатора является кодом-описанием объекта, и при этом кодирует черты, позволяющие дискриминатору отличать сгенерированные объекты от реальных, то можно заменить какую-то простую метрику (например, попиксельную), на метрику над активациями нейронов в каком-то из слоев:\n",
    "\n",
    "$$\n",
    "L(X_s, \\tilde X_s) \\longrightarrow L_d(d_l(X_s), d_l(\\tilde X_s)) \\\\\n",
    "\\tilde X_s = f_d(Z; \\theta_d), \\\\\n",
    "Z \\sim Q(X_s; \\theta_e),\n",
    "$$\n",
    "\n",
    "где $d_l$ - активации на $l$-ом слое дискриминатора, а $Q, \\ f_d$ - энкодер и декодер.\n",
    "\n",
    "При этом можно надеяться, что новая метрика $L_d$ будет лучше."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ниже приведена схема работы получившейся *VAE+GAN* сети, предлагаемая авторами ***[1]***.  \n",
    "Иллюстрация из ***[1]***  \n",
    "\n",
    "![](figs6/VAE_GAN_losses.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Здесь:\n",
    "\n",
    "- $X$ - входной объект из $P(X)$,\n",
    "- $Z_p$ - сэмплированный $Z$ из $P(Z)$,\n",
    "- $X_p$ - объект сгенерированный декодером из $Z_p$,\n",
    "- $\\tilde X$ - объект восстановленный из $X$,\n",
    "- $\\mathcal L_{prior} = KL \\left[ Q(Z|X)||P(Z) \\right]$ - лосс, заставляющий энкодер переводить $P(X)$ в нужное нам $P(Z)$ (точно как в части 3 про *VAE*),\n",
    "- $\\mathcal L_{llike}^{Dis_l} = L_d(d_l(X), d_l(\\tilde X))$ - метрика между активациями $l$-ого слоя дискриминатора $D$ на реальном $X$ и восстановленным $\\tilde X = f_d(Q(X))$,\n",
    "- $\\mathcal L_{GAN} = \\log(D(X)) + \\log(1 - D(f_d(Z))) + \\log(1 - D(f_d(Q(X))))$ - кросс-энтропия между реальным распределением лейблов настоящих/сгенерированных объектов, и распределением вероятности предсказываемым дискриминатором."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Как и в случае с *GAN* мы не можем обучать все 3 части сети одновременно. Дискриминатор надо обучать отдельно, в частности не нужно, чтобы дискриминатор пытался уменьшать $\\mathcal L_{llike}^{Dis_l}$, так как это схлопнет разницу активаций в 0. Поэтому обучение всех сетей надо ограничить только на релевантные им лоссы.  \n",
    "\n",
    "Схема предлагаемая авторами:\n",
    "\n",
    "$$\n",
    "\\theta_{Enc} = \\theta_{Enc} - \\Delta_{\\theta_{Enc}} (\\mathcal L_{prior} + \\mathcal L^{Dis_l}_{llike}), \\\\\n",
    "\\theta_{Dec} = \\theta_{Dec} - \\Delta_{\\theta_{Dec}} (\\gamma \\mathcal L^{Dis_l}_{llike} - \\mathcal L_{GAN}), \\\\\n",
    "\\theta_{Dis} = \\theta_{Dis} - \\Delta_{\\theta_{Dis}} (\\mathcal L_{GAN})\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Выше видно на каких лоссах какие сети учатся. Особое внимание разве что стоит уделить декодеру: он с одной стороны пытается уменьшить расстояние между входом и выходом в метрике *l-го* слоя дискриминатора ($\\mathcal L^{Dis_l}_{llike}$), а с другой пытается обмануть дискриминатор (увеличивая $\\mathcal L_{GAN}$). В статье авторы утверждают, что меняя коэффициент $\\gamma$ можно влиять на то, что важнее для сети: контент ($\\mathcal L^{Dis_l}_{llike}$) или стиль ($\\mathcal L_{GAN}$). Не могу, однако, сказать, что наблюдал этот эффект."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Код"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Код во многом повторяет то, что было в прошлых частях про чистые ***VAE*** и ***GAN***.\n",
    "\n",
    "Опять же сразу будем писать **conditional** модель."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Почти полное повторение того, что было в части про *GAN*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "\n",
    "from keras.layers import Dropout, BatchNormalization, Reshape, Flatten, RepeatVector\n",
    "from keras.layers import Lambda, Dense, Input, Conv2D, MaxPool2D, UpSampling2D, concatenate\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras.layers import Activation\n",
    "from keras.models import Model, load_model\n",
    "\n",
    "\n",
    "# Регистрация сессии в keras\n",
    "from keras import backend as K\n",
    "import tensorflow as tf\n",
    "sess = tf.Session()\n",
    "K.set_session(sess)\n",
    "\n",
    "\n",
    "# Импорт датаеста\n",
    "from keras.datasets import mnist\n",
    "from keras.utils import to_categorical\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "\n",
    "x_train = x_train.astype('float32') / 255.\n",
    "x_test  = x_test .astype('float32') / 255.\n",
    "x_train = np.reshape(x_train, (len(x_train), 28, 28, 1))\n",
    "x_test  = np.reshape(x_test,  (len(x_test),  28, 28, 1))\n",
    "\n",
    "y_train_cat = to_categorical(y_train).astype(np.float32)\n",
    "y_test_cat  = to_categorical(y_test).astype(np.float32)\n",
    "\n",
    "\n",
    "# Глобальные константы\n",
    "batch_size = 64\n",
    "batch_shape = (batch_size, 28, 28, 1)\n",
    "latent_dim = 8\n",
    "num_classes = 10\n",
    "dropout_rate = 0.3\n",
    "gamma = 1 # Коэффициент гамма\n",
    "\n",
    "\n",
    "# Итераторы тренировочных и тестовых батчей\n",
    "def gen_batch(x, y):\n",
    "    n_batches = x.shape[0] // batch_size\n",
    "    while(True):\n",
    "        idxs = np.random.permutation(y.shape[0])\n",
    "        x = x[idxs]\n",
    "        y = y[idxs]\n",
    "        for i in range(n_batches):\n",
    "            yield x[batch_size*i: batch_size*(i+1)], y[batch_size*i: batch_size*(i+1)]\n",
    "\n",
    "train_batches_it = gen_batch(x_train, y_train_cat)\n",
    "test_batches_it  = gen_batch(x_test,  y_test_cat)\n",
    "\n",
    "\n",
    "# Входные плейсхолдеры\n",
    "x_ = tf.placeholder(tf.float32, shape=(None, 28, 28, 1),  name='image')\n",
    "y_ = tf.placeholder(tf.float32, shape=(None, 10),         name='labels')\n",
    "z_ = tf.placeholder(tf.float32, shape=(None, latent_dim), name='z')\n",
    "\n",
    "img = Input(tensor=x_)\n",
    "lbl = Input(tensor=y_)\n",
    "z   = Input(tensor=z_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Описание моделей, от *GAN* отличается только добавленным энкодером"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def add_units_to_conv2d(conv2, units):\n",
    "    dim1 = int(conv2.shape[1])\n",
    "    dim2 = int(conv2.shape[2])\n",
    "    dimc = int(units.shape[1])\n",
    "    repeat_n = dim1*dim2\n",
    "    units_repeat = RepeatVector(repeat_n)(lbl)\n",
    "    units_repeat = Reshape((dim1, dim2, dimc))(units_repeat)\n",
    "    return concatenate([conv2, units_repeat])\n",
    "\n",
    "\n",
    "def apply_bn_relu_and_dropout(x, bn=False, relu=True, dropout=True):\n",
    "    if bn:\n",
    "        x = BatchNormalization(momentum=0.99, scale=False)(x)\n",
    "    if relu:\n",
    "        x = LeakyReLU()(x)\n",
    "    if dropout:\n",
    "        x = Dropout(dropout_rate)(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "with tf.variable_scope('encoder'):\n",
    "    x = Conv2D(32, kernel_size=(3, 3), strides=(2, 2), padding='same')(img)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "    x = MaxPool2D((2, 2), padding='same')(x)\n",
    "#     x = add_units_to_conv2d(x, lbl)\n",
    "\n",
    "    x = Conv2D(64, kernel_size=(3, 3), padding='same')(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "\n",
    "    x = Flatten()(x)\n",
    "    x = concatenate([x, lbl])\n",
    "    \n",
    "    h = Dense(64)(x)\n",
    "    h = apply_bn_relu_and_dropout(h)\n",
    "\n",
    "    z_mean    = Dense(latent_dim)(h)\n",
    "    z_log_var = Dense(latent_dim)(h)\n",
    "\n",
    "    def sampling(args):\n",
    "        z_mean, z_log_var = args\n",
    "        epsilon = K.random_normal(shape=(batch_size, latent_dim), mean=0., stddev=1.0)\n",
    "        return z_mean + K.exp(K.clip(z_log_var/2, -2, 2)) * epsilon\n",
    "    l = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_var])\n",
    "encoder = Model([img, lbl], [z_mean, z_log_var, l], name='Encoder')\n",
    "\n",
    "\n",
    "with tf.variable_scope('decoder'):\n",
    "    x = concatenate([z, lbl])\n",
    "    x = Dense(7*7*128)(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "    x = Reshape((7, 7, 128))(x)\n",
    "    x = UpSampling2D(size=(2, 2))(x)\n",
    "\n",
    "    x = Conv2D(64, kernel_size=(5, 5), padding='same')(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "\n",
    "    x = Conv2D(32, kernel_size=(3, 3), padding='same')(x)\n",
    "    x = UpSampling2D(size=(2, 2))(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "#     x = add_units_to_conv2d(x, lbl)\n",
    "    \n",
    "    decoded = Conv2D(1, kernel_size=(5, 5), activation='sigmoid', padding='same')(x)\n",
    "decoder = Model([z, lbl], decoded, name='Decoder')\n",
    "\n",
    "\n",
    "with tf.variable_scope('discrim'):\n",
    "    x = Conv2D(128, kernel_size=(7, 7), strides=(2, 2), padding='same')(img)\n",
    "    x = MaxPool2D((2, 2), padding='same')(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "    x = add_units_to_conv2d(x, lbl)\n",
    "\n",
    "    x = Conv2D(64, kernel_size=(3, 3), padding='same')(x)\n",
    "    x = MaxPool2D((2, 2), padding='same')(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "\n",
    "    # l-слой на котором будем сравнивать активации\n",
    "    l = Conv2D(16, kernel_size=(3, 3), padding='same')(x)\n",
    "    x = apply_bn_relu_and_dropout(x)\n",
    "\n",
    "    h = Flatten()(x)\n",
    "    d = Dense(1, activation='sigmoid')(h)\n",
    "discrim = Model([img, lbl], [d, l], name='Discriminator')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Построение графа вычислений на основе моделей"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "z_mean, z_log_var, encoded_img = encoder([img, lbl])\n",
    "\n",
    "decoded_img = decoder([encoded_img, lbl])\n",
    "decoded_z   = decoder([z,           lbl])\n",
    "\n",
    "discr_img,     discr_l_img     = discrim([img,         lbl])\n",
    "discr_dec_img, discr_l_dec_img = discrim([decoded_img, lbl])\n",
    "discr_dec_z,   discr_l_dec_z   = discrim([decoded_z,   lbl])\n",
    "\n",
    "cvae_model = Model([img, lbl], decoder([encoded_img, lbl]), name='cvae')\n",
    "cvae =  cvae_model([img, lbl])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Определение лоссов  \n",
    "Интересно, что получался лучше результат, если в качестве метрики на активациях слоев брать не *MSE*, а кросс-энтропию"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Базовые лоссы\n",
    "L_prior = -0.5*tf.reduce_sum(1. + tf.clip_by_value(z_log_var, -2, 2) - tf.square(z_mean) - tf.exp(tf.clip_by_value(z_log_var, -2, 2)))/28/28\n",
    "\n",
    "log_dis_img     = tf.log(discr_img + 1e-10)\n",
    "log_dis_dec_z   = tf.log(1. - discr_dec_z + 1e-10)\n",
    "log_dis_dec_img = tf.log(1. - discr_dec_img + 1e-10)\n",
    "\n",
    "L_GAN = -1/4*tf.reduce_sum(log_dis_img + 2*log_dis_dec_z + log_dis_dec_img)/28/28\n",
    "\n",
    "# L_dis_llike = tf.reduce_sum(tf.square(discr_l_img - discr_l_dec_img))/28/28\n",
    "L_dis_llike = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.sigmoid(discr_l_img),\n",
    "                                                                    logits=discr_l_dec_img))/28/28\n",
    "\n",
    "\n",
    "# Лоссы энкодера, декодера, дискриминатора\n",
    "L_enc = L_dis_llike + L_prior \n",
    "L_dec = gamma * L_dis_llike - L_GAN\n",
    "L_dis = L_GAN\n",
    "\n",
    "\n",
    "# Определение шагов оптимизатора\n",
    "optimizer_enc = tf.train.RMSPropOptimizer(0.001)\n",
    "optimizer_dec = tf.train.RMSPropOptimizer(0.0003)\n",
    "optimizer_dis = tf.train.RMSPropOptimizer(0.001)\n",
    "\n",
    "encoder_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, \"encoder\")\n",
    "decoder_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, \"decoder\")\n",
    "discrim_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, \"discrim\")\n",
    "\n",
    "step_enc = optimizer_enc.minimize(L_enc, var_list=encoder_vars)\n",
    "step_dec = optimizer_dec.minimize(L_dec, var_list=decoder_vars)\n",
    "step_dis = optimizer_dis.minimize(L_dis, var_list=discrim_vars)\n",
    "\n",
    "\n",
    "def step(image, label, zp):\n",
    "    l_prior, dec_image, l_dis_llike, l_gan, _, _ = sess.run([L_prior, decoded_z, L_dis_llike, L_GAN, step_enc, step_dec],\n",
    "                                                            feed_dict={z:zp, img:image, lbl:label, K.learning_phase():1})\n",
    "    return l_prior, dec_image, l_dis_llike, l_gan\n",
    "\n",
    "def step_d(image, label, zp):\n",
    "    l_gan, _ = sess.run([L_GAN, step_dis], feed_dict={z:zp, img:image, lbl:label, K.learning_phase():1})\n",
    "    return l_gan\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Функции рисования картинок после и в процессе тренировки (скрыто)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "digit_size = 28\n",
    "def plot_digits(*args, invert_colors=False):\n",
    "    args = [x.squeeze() for x in args]\n",
    "    n = min([x.shape[0] for x in args])\n",
    "    figure = np.zeros((digit_size * len(args), digit_size * n))\n",
    "\n",
    "    for i in range(n):\n",
    "        for j in range(len(args)):\n",
    "            figure[j * digit_size: (j + 1) * digit_size,\n",
    "                   i * digit_size: (i + 1) * digit_size] = args[j][i].squeeze()\n",
    "\n",
    "    if invert_colors:\n",
    "        figure = 1-figure\n",
    "\n",
    "    plt.figure(figsize=(2*n, 2*len(args)))\n",
    "    plt.imshow(figure, cmap='Greys_r')\n",
    "    plt.grid(False)\n",
    "    ax = plt.gca()\n",
    "    ax.get_xaxis().set_visible(False)\n",
    "    ax.get_yaxis().set_visible(False)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# Массивы в которые будем сохранять результаты, для последующей визуализации\n",
    "figs = [[] for x in range(num_classes)]\n",
    "periods = []\n",
    "\n",
    "save_periods = list(range(100)) + list(range(100, 1000, 10))\n",
    "\n",
    "n = 15 # Картинка с 15x15 цифр\n",
    "from scipy.stats import norm\n",
    "# Так как сэмплируем из N(0, I), то сетку узлов, в которых генерируем цифры берем из обратной функции распределения\n",
    "grid_x = norm.ppf(np.linspace(0.05, 0.95, n))\n",
    "grid_y = norm.ppf(np.linspace(0.05, 0.95, n))\n",
    "grid_y = norm.ppf(np.linspace(0.05, 0.95, n))\n",
    "def draw_manifold(label, show=True):\n",
    "    # Рисование цифр из многообразия\n",
    "    figure = np.zeros((digit_size * n, digit_size * n))\n",
    "    input_lbl = np.zeros((1, 10))\n",
    "    input_lbl[0, label] = 1\n",
    "    for i, yi in enumerate(grid_x):\n",
    "        for j, xi in enumerate(grid_y):\n",
    "            z_sample = np.zeros((1, latent_dim))\n",
    "            z_sample[:, :2] = np.array([[xi, yi]])\n",
    "\n",
    "            x_decoded = sess.run(decoded_z, feed_dict={z:z_sample, lbl:input_lbl, K.learning_phase():0})\n",
    "            digit = x_decoded[0].squeeze()\n",
    "            figure[i * digit_size: (i + 1) * digit_size,\n",
    "                   j * digit_size: (j + 1) * digit_size] = digit\n",
    "    if show:\n",
    "        # Визуализация\n",
    "        plt.figure(figsize=(15, 15))\n",
    "        plt.imshow(figure, cmap='Greys')\n",
    "        plt.grid(False)\n",
    "        ax = plt.gca()\n",
    "        ax.get_xaxis().set_visible(False)\n",
    "        ax.get_yaxis().set_visible(False)\n",
    "        plt.show()\n",
    "    return figure\n",
    "\n",
    "\n",
    "# Рисование распределения z\n",
    "def draw_z_distr(z_predicted):\n",
    "    im = plt.scatter(z_predicted[:, 0], z_predicted[:, 1])\n",
    "    im.axes.set_xlim(-5, 5)\n",
    "    im.axes.set_ylim(-5, 5)\n",
    "    plt.show()\n",
    "    \n",
    "\n",
    "def on_n_period(period):\n",
    "    n_compare = 10\n",
    "\n",
    "    clear_output() # Не захламляем output\n",
    "\n",
    "    # Сравнение реальных и декодированных цифр\n",
    "    b = next(test_batches_it)\n",
    "    decoded = sess.run(cvae, feed_dict={img:b[0], lbl:b[1], K.learning_phase():0})\n",
    "    plot_digits(b[0][:n_compare], decoded[:n_compare])\n",
    "\n",
    "    # Рисование многообразия для рандомного y\n",
    "    draw_lbl = np.random.randint(0, num_classes)    \n",
    "    print(draw_lbl)\n",
    "    for label in range(num_classes):\n",
    "        figs[label].append(draw_manifold(label, show=label==draw_lbl))\n",
    "\n",
    "    xs = x_test[y_test == draw_lbl]\n",
    "    ys = y_test_cat[y_test == draw_lbl]\n",
    "    z_predicted = sess.run(z_mean, feed_dict={img:xs, lbl:ys, K.learning_phase():0})\n",
    "    draw_z_distr(z_predicted)\n",
    "    \n",
    "    periods.append(period)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Процесс обучения"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "nb_step = 3 # Количество шагов во внутреннем цикле\n",
    "\n",
    "\n",
    "batches_per_period = 3\n",
    "for i in range(48000):\n",
    "    print('.', end='')\n",
    "\n",
    "    # Шаги обучения дискриминатора\n",
    "    for j in range(nb_step):\n",
    "        b0, b1 = next(train_batches_it)\n",
    "        zp = np.random.randn(batch_size, latent_dim)\n",
    "        l_g = step_d(b0, b1, zp)\n",
    "        if l_g < 1.0:\n",
    "            break\n",
    "        \n",
    "    # Шаг обучения декодера и энкодера\n",
    "    for j in range(nb_step):\n",
    "        l_p, zx, l_d, l_g = step(b0, b1, zp)\n",
    "        if l_g > 0.4:\n",
    "            break\n",
    "        b0, b1 = next(train_batches_it)\n",
    "        zp = np.random.randn(batch_size, latent_dim)\n",
    "\n",
    "    # Периодическая визуализация результата\n",
    "    if not i % batches_per_period:\n",
    "        period = i // batches_per_period\n",
    "        if period in save_periods:\n",
    "            on_n_period(period)\n",
    "        print(i, l_p, l_d, l_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Функция рисования гифок (скрыто)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from matplotlib.animation import FuncAnimation\n",
    "from matplotlib import cm\n",
    "import matplotlib\n",
    "\n",
    "\n",
    "def make_2d_figs_gif(figs, periods, c, fname, fig, batches_per_period): \n",
    "    norm = matplotlib.colors.Normalize(vmin=0, vmax=1, clip=False)\n",
    "    im = plt.imshow(np.zeros((28,28)), cmap='Greys', norm=norm)\n",
    "    plt.grid(None)\n",
    "    plt.title(\"Label: {}\\nBatch: {}\".format(c, 0))\n",
    "\n",
    "    def update(i):\n",
    "        im.set_array(figs[i])\n",
    "        im.axes.set_title(\"Label: {}\\nBatch: {}\".format(c, periods[i]*batches_per_period))\n",
    "        im.axes.get_xaxis().set_visible(False)\n",
    "        im.axes.get_yaxis().set_visible(False)\n",
    "        return im\n",
    "    \n",
    "    anim = FuncAnimation(fig, update, frames=range(len(figs)), interval=100)\n",
    "    anim.save(fname, dpi=80, writer='ffmpeg')\n",
    "\n",
    "for label in range(num_classes):\n",
    "    make_2d_figs_gif(figs[label], periods, label, \"./figs6/manifold_{}.mp4\".format(label), plt.figure(figsize=(10,10)), batches_per_period)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Так как у нас снова модель на основе автоэнкодера, мы можем применять перенос стиля"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Трансфер стиля\n",
    "def style_transfer(X, lbl_in, lbl_out):\n",
    "    rows = X.shape[0]\n",
    "    if isinstance(lbl_in, int):\n",
    "        label = lbl_in\n",
    "        lbl_in = np.zeros((rows, 10))\n",
    "        lbl_in[:, label] = 1\n",
    "    if isinstance(lbl_out, int):\n",
    "        label = lbl_out\n",
    "        lbl_out = np.zeros((rows, 10))\n",
    "        lbl_out[:, label] = 1\n",
    "    # Кодирем стиль входящего изображения\n",
    "    zp = sess.run(z_mean, feed_dict={img:X, lbl:lbl_in, K.learning_phase():0})\n",
    "    # Восстанавливаем из этого стиля, заменяя лейбл\n",
    "    created = sess.run(decoded_z, feed_dict={z:zp, lbl:lbl_out, K.learning_phase():0})\n",
    "    return created\n",
    "\n",
    "\n",
    "# Картинка трансфера стиля\n",
    "def draw_random_style_transfer(label):\n",
    "    n = 10\n",
    "    generated = []\n",
    "    idxs = np.random.permutation(y_test.shape[0])\n",
    "    x_test_permut = x_test[idxs]\n",
    "    y_test_permut = y_test[idxs]\n",
    "    prot = x_test_permut[y_test_permut == label][:batch_size]\n",
    "\n",
    "    for i in range(num_classes):\n",
    "        generated.append(style_transfer(prot, label, i)[:n])\n",
    "\n",
    "    generated[label] = prot\n",
    "\n",
    "    plot_digits(*generated, invert_colors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "draw_random_style_transfer(7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Результаты"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Сравнение с простым CVAE\n",
    "Сверху оригиналы цифр, снизу восстановленные\n",
    "\n",
    "\n",
    "*CVAE*, скрытая размерность - 2\n",
    "![](figs4/map.png)\n",
    "\n",
    "*CVAE+GAN*, скрытая размерность - 2\n",
    "![](figs6/image_comp_VAE_GAN_2.png)\n",
    "\n",
    "*CVAE+GAN*, скрытая размерность - 8\n",
    "![](figs6/image_comp_VAE_GAN_8.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Сгенерированные цифры каждого лейбла сэмплированные из $N(0|I)$\n",
    "<img src=\"./figs6/manifold_0.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_1.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_2.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_3.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_4.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_5.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_6.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_7.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_8.png\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_9.png\" width=\"600\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Процесс обучения (тяжелые гифки) (скрыто)\n",
    "Сгенерированные цифры каждого лейбла сэмплированные из $N(0|I)$\n",
    "<img src=\"./figs6/manifold_0.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_1.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_2.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_3.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_4.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_5.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_6.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_7.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_8.gif\" width=\"600\"/>\n",
    "<img src=\"./figs6/manifold_9.gif\" width=\"600\"/>\n",
    "\n",
    "### Трансфер стиля\n",
    "За основу брались \"7\", из стиля которых создавались уже остальные цифры (здесь $\\dim Z$ = 8)\n",
    "\n",
    "Вот так было с простым *CVAE*:\n",
    "![](figs4/style_transfer.png)\n",
    "\n",
    "А вот так стало:\n",
    "![](figs6/style_transfer.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Заключение\n",
    "На мой взгляд получилось очень неплохо. Пройдя путь от простейших *автоэнкодеров* мы дошли до генеративных моделей, а именно до *VAE*, *GAN*, поняли что такое *conditional* модели, и почему важна метрика.  \n",
    "Мы также научились пользоваться *keras'ом* и совмещать его с голым *tensorflow*.\n",
    "\n",
    "Всем спасибо за внимание, надеюсь было интересно!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Полезные ссылки и литература\n",
    "\n",
    "Оригинальная статья:  \n",
    "[1] Autoencoding beyond pixels using a learned similarity metric, Larsen et al, 2016, https://arxiv.org/abs/1512.09300\n",
    "\n",
    "Туториал по *VAE*:  \n",
    "[2] Tutorial on Variational Autoencoders, Carl Doersch, 2016, https://arxiv.org/abs/1606.05908\n",
    "\n",
    "Туториал по использованию *keras* вместе с *tensorflow*:  \n",
    "[3] https://blog.keras.io/keras-as-a-simplified-interface-to-tensorflow-tutorial.html"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}