VAE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Автоэнкодеры в Keras\n",
    "\n",
    "# Часть 3: Вариационные автоэнкодеры (VAE)"
   ]
  },
  {
   "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": [
    "В прошлой части мы уже обсуждали, что такое скрытые переменные, взглянули на их распределение, а так же поняли, что из распределения скрытых переменных в обычных автоэнкодерах сложно генерировать новые объекты. Для того, чтобы можно было генерировать новые объекты, пространство *скрытых переменных* (*latent variables*) должно быть предсказуемым. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Имея какое-то одно распределение $Z$ можно получить произвольное другое $X = g(Z)$, например,\n",
    "\n",
    "пусть $Z$ - обычное нормальное распределение, $g(Z) = \\frac{Z}{|Z|}+ \\frac{Z}{10}$ - тоже случайное распределение, но выглядит совсем по другому"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Код: (скрыто)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "\n",
    "\n",
    "Z = np.random.randn(150, 2)\n",
    "X = Z/(np.sqrt(np.sum(Z*Z, axis=1))[:, None]) + Z/10\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, sharex=False, figsize=(16,8))\n",
    "\n",
    "ax = axs[0]\n",
    "ax.scatter(Z[:,0], Z[:,1])\n",
    "ax.grid(True)\n",
    "ax.set_xlim(-5, 5)\n",
    "ax.set_ylim(-5, 5)\n",
    "\n",
    "ax = axs[1]\n",
    "ax.scatter(X[:,0], X[:,1])\n",
    "ax.grid(True)\n",
    "ax.set_xlim(-2, 2)\n",
    "ax.set_ylim(-2, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](figs3/custom_dist_f.png)\n",
    "Пример выше из ***[1]***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Таким образом, если подобрать правильные функции, то можно отобразить пространства скрытых переменных обычных автоэнкодеров в какие-то хорошие пространства, например такие, где распределение нормально. А потом обратно.\n",
    "\n",
    "С другой стороны специально учиться отображать одни скрытые пространства в другие вовсе не обязательно. Если есть какие-то полезные скрытые пространства, то правильный автоэнкодер научится им по пути сам, но отображать, в конечно итоге, будет в нужное нам пространство."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***Вариационные автоэнкодеры*** (*Variational Autoencoders*) - это автоэнкодеры, которые учатся отображать объекты в заданное скрытое пространство и, соответственно, сэмплить из него. Поэтому *вариационные автоэнкодеры* относят так же к семейству генеративных моделей."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](figs3/vae_diagram.png)\n",
    "Иллюстрация из ***[2]***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Далее выжимка теории из ***[1]*** лежащая в основе *VAE*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Пусть $Z$ - скрытые переменные, а $X$ - данные. \n",
    "На примере нарисованых цифр рассмотрим естесственный генеративный процесс, который сгенерировал нашу выборку:\n",
    "$$\n",
    "P(X) = \\int_{z} P(X|Z)P(Z)dZ\n",
    "$$\n",
    "\n",
    "- $P(X)$ вероятностное распределение изображений цифр на картинках, т.е. вероятность конкретного изображения цифры впринципе быть нарисованым (если картинка не похожа на цифру, то эта вероятность крайне мала, и наоборот),\n",
    "- $P(Z)$ - вероятностное распределение скрытых факторов, например, распределение толщины штриха,\n",
    "- $P(X|Z)$ - распределение вероятности картинок при заданных скрытых факторах, одни и те же факторы могут привезти к разным картинкам (один и тот же человек в одних и тех же условиях не рисует абсолютно одинаковые цифры)\n",
    "\n",
    "Представим $P(X|Z)$ как сумму некоторой генерирующей функции $f(Z)$ и некоторого сложного шума $\\epsilon$\n",
    "\n",
    "$$\n",
    "P(X|Z) = f(Z) + \\epsilon\n",
    "$$\n",
    "\n",
    "Мы хотим построить некоторый искусственный генеративный процесс, который будет создавать объекты близкие в некоторой метрике к тренировачным $X$.\n",
    "\n",
    "$$\n",
    "P(X;\\theta) = \\int_{z} P(X|Z;\\theta)P(Z)dZ \\ \\ \\ (1)\n",
    "$$\n",
    "и снова\n",
    "$$\n",
    "P(X|Z;\\theta) = f(Z;\\theta) + \\epsilon\n",
    "$$\n",
    "\n",
    "$f(Z;\\theta)$ - некоторое семейсто функций, которое представляет наша модель, а $\\theta$ - ее параметры. Выбирая метрику - мы выбираем то, какого вида нам представляется шум $\\epsilon$. Если метрика $L_2$, то мы считаем шум нормальным и тогда:\n",
    "\n",
    "$$\n",
    "P(X|Z;\\theta) = N(X|f(Z;\\theta), \\sigma^2 I),\n",
    "$$\n",
    "\n",
    "По принципу максимального правдоподобия нам остается оптимизировать параметры $\\theta$ для того, чтобы максимизировать $P(X)$, т.е. вероятность появления объектов из выборки."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проблема в том, что оптимизировтаь интеграл (1) напрямую мы не можем: пространство может быть высокоразмерное, объектов много, да и метрика плохая. С другой стороны, если задуматься, то к каждому конкретному $X$ может привезти лишь очень небольшое подмножество $Z$, для остальных же $P(X|Z)$ будет очень близок к нулю. \n",
    "И при оптимизации достаточно сэмплить только из хороших $Z$.\n",
    "\n",
    "Для того чтобы знать из каких $Z$ нам надо сэмплить, введем новое распределение $Q(Z|X)$, которое в зависимости от $X$ будет показывать распределение $Z \\sim Q$, которое могло привезти к этому $X$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Запишем сперва расстояние Кульбака-Лейблера (несимметричная мера \"похожести\" двух распределений, подробнее ***[3]***) между\n",
    "$Q(Z|X)$ и реальным $P(Z|X)$:\n",
    "\n",
    "$$\n",
    "KL[Q(Z|X)||P(Z|X)] = \\mathbb{E}_{Z \\sim Q}[\\log Q(Z|X) - \\log P(Z|X)]\n",
    "$$\n",
    "\n",
    "Применяем формулу Байеса:\n",
    "\n",
    "$$\n",
    "KL[Q(Z|X)||P(Z|X)] = \\mathbb{E}_{Z \\sim Q}[\\log Q(Z|X) - \\log P(X|Z) - \\log P(Z)] + \\log P(X)\n",
    "$$\n",
    "\n",
    "Выделяем еще одно расстояние Кульбака-Лейблера:\n",
    "\n",
    "$$\n",
    "KL[Q(Z|X)||P(Z|X)] = KL[Q(Z|X)||\\log P(Z)] - \\mathbb{E}_{Z \\sim Q}[\\log P(X|Z)] + \\log P(X)\n",
    "$$\n",
    "\n",
    "В итоге получаем тождество:\n",
    "\n",
    "$$\n",
    "\\log P(X) - KL[Q(Z|X)||P(Z|X)] = \\mathbb{E}_{Z \\sim Q}[\\log P(X|Z)] - KL[Q(Z|X)||P(Z)]\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Это тождество - краеугольный камень *вариационных автоэнкодеров*, оно верно для любых $Q(Z|X)$ и $P(X,Z)$.\n",
    "\n",
    "Пусть $Q(Z|X)$ и $P(X|Z)$ зависят от параметров: $Q(Z|X;\\theta_1)$ и $P(X|Z;\\theta_2)$, а $P(Z)$ - нормальное $N(0,I)$, тогда получаем:\n",
    "\n",
    "$$\n",
    "\\log P(X;\\theta_2) - KL[Q(Z|X;\\theta_1)||P(Z|X;\\theta_2)] = \\mathbb{E}_{Z \\sim Q}[\\log P(X|Z;\\theta_2)] - KL[Q(Z|X;\\theta_1)||N(0,I)]\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Взглянем повнимательнее на то, что у нас получилось:\n",
    "- во-первых, $Q(Z|X;\\theta_1)$, $P(X|Z;\\theta_2)$ подозрительно похожи на энкодер и декодер (точнее декодер это $f$ в выражении $P(X|Z;\\theta_2) = f(Z;\\theta_2) + \\epsilon$)\n",
    "- слева в тождестве - значение, которое мы хотим максимизировать для элементов нашей тренировачной выборки $X$ + некоторая ошибка $KL$ ($KL(x,y) \\ge 0 \\ \\ \\forall x,y$), которая, будем надеяться, при достаточной емкости $Q$ уйдет в 0,\n",
    "- справа значение, которое мы можем оптимизировать градиентным спуском, где первый член имеет смысл качества предсказания $X$ декодером по значениям $Z$, а второй член, это расстояние К-Л между распределением $Z \\sim Q$, которое предсказывает энкодер для конкретного $X$, и распределением $Z$ для всех $X$ сразу"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для того, чтобы иметь возможность оптимизировать правую часть градиентным спуском, осталось разобраться с двумя вещами:\n",
    "#### 1. Точнее определим что такое $Q(Z|X;\\theta_1)$\n",
    "Обычно $Q$ выбирается нормальным распределением:\n",
    "\n",
    "$$\n",
    "Q(Z|X;\\theta_1) = N(\\mu(X;\\theta_1), \\Sigma(X;\\theta_1))\n",
    "$$\n",
    "То есть энкодер для каждого $X$ предсказывает 2 значения: среднее $\\mu$ и вариацию $\\Sigma$ нормального распределения, из которого уже сэмплируются значения. Работает это все примерно вот так:  \n",
    "![](figs3/encoder_vae_diagram.png)\n",
    "Иллюстрация из ***[2]***\n",
    "\n",
    "\n",
    "При том, что для каждой отдельной точки данных $X$ энкодер предсказывает некоторое нормальное распределение $P(Z|X) = N(\\mu(X), \\Sigma(X))$, для априорного распределения $Х$: $P(Z) = N(0, I)$, что получается из формулы, и это потрясающе.\n",
    "![](figs3/kl_divergence_diagram.png)\n",
    "Иллюстрация из ***[2]***\n",
    "\n",
    "\n",
    "При этом $KL[Q(Z|X;\\theta_1)||N(0,I)]$ принимает вид:\n",
    "\n",
    "$$\n",
    "KL[Q(Z|X;\\theta_1)||N(0,I)] = \\frac{1}{2}\\left(tr(\\Sigma(X)) + \\mu(X)^T\\mu(X) - k - \\log \\det \\Sigma(X) \\right)\n",
    "$$\n",
    "\n",
    "#### 2. Разберемся с тем, как распространять ошибки через $\\mathbb{E}_{Z \\sim Q}[\\log P(X|Z;\\theta_2)]$\n",
    "Дело в том, что здесь мы берем случайные значения $Z \\sim Q(Z|X;\\theta_1)$ и передаем их в декодер.\n",
    "Ясно, что распросранять ошибки через случайные значения напрямую нельзя, поэтому используется так называемый *трюк с репараметризацией* (*reparametrization trick*).\n",
    "\n",
    "Схема получается вот такая:\n",
    "![](figs3/reparam_trick.png)\n",
    "Иллюстрация из ***[1]***\n",
    "\n",
    "Здесь на левой картинке схема без трюка, а на правой с трюком.\n",
    "Красным цветом показано семплирование, а синим вычисление ошибки.  \n",
    "То есть по сути просто берем предсказанное энкодером стандартное отклонение $\\Sigma$ умножаем на случайное число из $N(0,I)$ и добавляем предсказанное среднее $\\mu$.\n",
    "\n",
    "Прямое растространение на обеих схемах абсолютно одинаковое, однако на правой схеме работает обратное распространение ошибки."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "После того, как мы обучили такой вариационный автоэнкодер, декодер становится полноправной генеративной моделью. По сути и энкодер то нужен в основном для того, чтобы обучить декодер отдельно быть генеративной моделью. \n",
    "![](figs3/vae_decoder_diagram.png)\n",
    "Иллюстрация из ***[2]***\n",
    "<img src=\"./figs3/generator.png\" width=\"300\"/>\n",
    "Иллюстрация из ***[1]***\n",
    "\n",
    "Но то, что энкодер и декодер вместо образуют еще и полноценный автоэнкодер, это очень приятный плюс."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VAE в Keras "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Теперь, когда мы разобрались в том, что такое вариационные автоэнкодеры, напишем такой на *Keras*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Импортируем необходимые библиотеки и датасет"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "\n",
    "from keras.datasets import mnist\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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим основные параметры. Скрытое пространство возьмем размерности 2, чтобы позже генерировать из него и визуализировать результат.  \n",
    "***Замечание***: размерность 2 крайне мала, особенно в метрике $L_2$, поэтому следует ожидать, что цифры получатся очень размытыми."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_size = 500\n",
    "latent_dim = 2\n",
    "dropout_rate = 0.3\n",
    "start_lr = 0.0001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Напишем модели вариационного автоэнкодера. \n",
    "\n",
    "Для того, чтобы обучение происходило быстрее и более качественно, добавим слои *dropout* и *batch normalization*.\n",
    "А в декодере используем в качестве активации *leaky ReLU*, которую добавляем отдельным слоем после *dense* слоев без активации.\n",
    "\n",
    "Функция *sampling* реализует сэмплирование значений $Z$ из $Q(X)$ с использованием трюка репараметризации.\n",
    "\n",
    "*vae_loss* это правая часть из уравнения:\n",
    "\n",
    "$$\n",
    "\\log P(X;\\theta_2) - KL[Q(Z|X;\\theta_1)||P(Z|X;\\theta_2)] = \\mathbb{E}_{Z \\sim Q}[\\log P(X|Z;\\theta_2)] - \\left(\\frac{1}{2}\\left(tr(\\Sigma(X)) + \\mu(X)^T\\mu(X) - k - \\log \\det \\Sigma(X) \\right)\\right)\n",
    "$$\n",
    "\n",
    "далее будет использоваться в качестве лосса."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from keras.layers import Input, Dense \n",
    "from keras.layers import BatchNormalization, Dropout, Flatten, Reshape, Lambda\n",
    "from keras.models import Model\n",
    "\n",
    "from keras.objectives import binary_crossentropy\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras import backend as K\n",
    "\n",
    "\n",
    "def create_vae():\n",
    "    models = {}\n",
    "\n",
    "    # Добавим Dropout и BatchNormalization\n",
    "    def apply_bn_and_dropout(x):\n",
    "        return Dropout(dropout_rate)(BatchNormalization()(x))\n",
    "\n",
    "    # Энкодер\n",
    "    input_img = Input(batch_shape=(batch_size, 28, 28, 1))\n",
    "    x = Flatten()(input_img)\n",
    "    x = Dense(256, activation='relu')(x)\n",
    "    x = apply_bn_and_dropout(x)\n",
    "    x = Dense(128, activation='relu')(x)\n",
    "    x = apply_bn_and_dropout(x)\n",
    "\n",
    "    # Предсказываем параметры распределений\n",
    "    # Вместо того, чтобы предсказывать стандартное отклонение, предсказываем логарифм вариации\n",
    "    z_mean = Dense(latent_dim)(x)\n",
    "    z_log_var = Dense(latent_dim)(x)\n",
    "\n",
    "    # Сэмплирование из Q с трюком репараметризации\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(z_log_var / 2) * epsilon\n",
    "    l = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_var])\n",
    "\n",
    "    models[\"encoder\"]  = Model(input_img, l, 'Encoder') \n",
    "    models[\"z_meaner\"] = Model(input_img, z_mean, 'Enc_z_mean')\n",
    "    models[\"z_lvarer\"] = Model(input_img, z_log_var, 'Enc_z_log_var')\n",
    "\n",
    "    # Декодер\n",
    "    z = Input(shape=(latent_dim, ))\n",
    "    x = Dense(128)(z)\n",
    "    x = LeakyReLU()(x)\n",
    "    x = apply_bn_and_dropout(x)\n",
    "    x = Dense(256)(x)\n",
    "    x = LeakyReLU()(x)\n",
    "    x = apply_bn_and_dropout(x)\n",
    "    x = Dense(28*28, activation='sigmoid')(x)\n",
    "    decoded = Reshape((28, 28, 1))(x)\n",
    "\n",
    "    models[\"decoder\"] = Model(z, decoded, name='Decoder')\n",
    "    models[\"vae\"]     = Model(input_img, models[\"decoder\"](models[\"encoder\"](input_img)), name=\"VAE\")\n",
    "\n",
    "    def vae_loss(x, decoded):\n",
    "        x = K.reshape(x, shape=(batch_size, 28*28))\n",
    "        decoded = K.reshape(decoded, shape=(batch_size, 28*28))\n",
    "        xent_loss = 28*28*binary_crossentropy(x, decoded)\n",
    "        kl_loss = -0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)\n",
    "        return (xent_loss + kl_loss)/2/28/28\n",
    "\n",
    "    return models, vae_loss\n",
    "\n",
    "models, vae_loss = create_vae()\n",
    "vae = models[\"vae\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***Замечание***: мы использовали *Lambda*-слой с функцией сэмплирующей из $N(0, I)$ из нижележащего фреймворка, которая явно требует размер батча. Во всех моделях, в которых присутствует этот слой мы теперь вынуждены передавать именно такой размер батча (то есть в \"encoder\" и \"vae*)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Функцией оптимизации возьмем *Adam* или *RMSprop*, обе показывают хорошие результаты.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from keras.optimizers import Adam, RMSprop\n",
    "\n",
    "vae.compile(optimizer=Adam(start_lr), loss=vae_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Код рисования рядов цифр и цифр из многообразия"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "digit_size = 28\n",
    "\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 = 15 # Картинка с 15x15 цифр\n",
    "digit_size = 28\n",
    "\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",
    "\n",
    "def draw_manifold(generator, show=True):\n",
    "    # Рисование цифр из многообразия\n",
    "    figure = np.zeros((digit_size * n, digit_size * n))\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 = generator.predict(z_sample)\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_r')\n",
    "        plt.grid(None)\n",
    "        ax = plt.gca()\n",
    "        ax.get_xaxis().set_visible(False)\n",
    "        ax.get_yaxis().set_visible(False)        \n",
    "        plt.show()\n",
    "    return figure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Часто в процессе обучения модели требуется выполнять какие-то действия: изменять *learning_rate*, сохранять промежуточные результаты, сохранять модель, рисовать картинки и т.д.\n",
    "\n",
    "Для этого в *keras* есть коллбэки, которые передаются в метод *fit* перед началом обучения. Например, чтобы влиять на *learning rate* в процессе обучения есть такие коллбэки, как *LearningRateScheduler*, *ReduceLROnPlateau*, чтобы сохранять модель - *ModelCheckpoint*.\n",
    "\n",
    "Отдельный коллбэк нужен для того, чтобы следить за процессом обучения в *TensorBoard*. Он автоматически будет добавлять в файл логов все метрики и лоссы, которые считаются между эпохами.\n",
    "\n",
    "Для случая, когда требуется выполнения произвольных функций в процессе обучения, существует *LambdaCallback*. Он запускает выполнение произвольных функций в заданные моменты обучения, например между эпохами или батчами.  \n",
    "Будем следить за процессом обучения, изучая, как генерируются цифры из $N(0, I)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "from keras.callbacks import LambdaCallback, ReduceLROnPlateau, TensorBoard\n",
    "\n",
    "# Массивы в которые будем сохранять результаты, для последующей визуализации\n",
    "figs = []\n",
    "latent_distrs = []\n",
    "epochs = []\n",
    "\n",
    "# Эпохи в которые будем сохранять\n",
    "save_epochs = set(list((np.arange(0, 59)**1.701).astype(np.int)) + list(range(10)))\n",
    "\n",
    "# Отслеживать будем на вот этих цифрах\n",
    "imgs = x_test[:batch_size]\n",
    "n_compare = 10\n",
    "\n",
    "# Модели\n",
    "generator      = models[\"decoder\"]\n",
    "encoder_mean   = models[\"z_meaner\"]\n",
    "\n",
    "\n",
    "# Фунция, которую будем запускать после каждой эпохи\n",
    "def on_epoch_end(epoch, logs):\n",
    "    if epoch in save_epochs:\n",
    "        clear_output() # Не захламляем output\n",
    "\n",
    "        # Сравнение реальных и декодированных цифр\n",
    "        decoded = vae.predict(imgs, batch_size=batch_size)\n",
    "        plot_digits(imgs[:n_compare], decoded[:n_compare])\n",
    "\n",
    "        # Рисование многообразия\n",
    "        figure = draw_manifold(generator, show=True)\n",
    "\n",
    "        # Сохранение многообразия и распределения z для создания анимации после\n",
    "        epochs.append(epoch)\n",
    "        figs.append(figure)\n",
    "        latent_distrs.append(encoder_mean.predict(x_test, batch_size))\n",
    "\n",
    "        \n",
    "# Коллбэки\n",
    "pltfig = LambdaCallback(on_epoch_end=on_epoch_end)\n",
    "# lr_red = ReduceLROnPlateau(factor=0.1, patience=25)\n",
    "tb     = TensorBoard(log_dir='./logs')\n",
    "\n",
    "\n",
    "# Запуск обучения \n",
    "vae.fit(x_train, x_train, shuffle=True, epochs=1000,\n",
    "        batch_size=batch_size,\n",
    "        validation_data=(x_test, x_test),\n",
    "        callbacks=[pltfig, tb],\n",
    "        verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Теперь, если установлен *TensorBoard*, можно следить за процессом обучения."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Вот как этот энкодер восстанавливает изображения:\n",
    "![](./figs3/image_comp.png)\n",
    "\n",
    "А вот результат сэмплирования из $N(0|I)$\n",
    "<img src=\"./figs3/manifold_all.png\" width=\"600\"/>\n",
    "\n",
    "Вот так выглядит процесс обучения генерации цифр (скрыто):\n",
    "<img src=\"./figs3/manifold.gif\" width=\"600\"/>\n",
    "Распределение кодов в скрытом пространстве (скрыто).  \n",
    "<img src=\"./figs3/z_distr.gif\" width=\"600\"/>\n",
    "\n",
    "Не идеально нормальное, но довольно близко (особенно, учитывая, что размерность скрытого пространства всего 2).\n",
    "\n",
    "Кривая обучения в *TensorBoard*\n",
    "<img src=\"./figs3/val_loss_tb.png\" width=\"800\"/>"
   ]
  },
  {
   "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",
    "def make_2d_figs_gif(figs, epochs, fname, fig): \n",
    "    norm = matplotlib.colors.Normalize(vmin=0, vmax=1, clip=False)\n",
    "    im = plt.imshow(np.zeros((28,28)), cmap='Greys_r', norm=norm)\n",
    "    plt.grid(None)\n",
    "    plt.title(\"Epoch: \" + str(epochs[0]))\n",
    "\n",
    "    def update(i):\n",
    "        im.set_array(figs[i])\n",
    "        im.axes.set_title(\"Epoch: \" + str(epochs[i]))\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='imagemagick')\n",
    "\n",
    "def make_2d_scatter_gif(zs, epochs, c, fname, fig):\n",
    "    im = plt.scatter(zs[0][:, 0], zs[0][:, 1], c=c, cmap=cm.coolwarm)\n",
    "    plt.colorbar()\n",
    "    plt.title(\"Epoch: \" + str(epochs[0]))\n",
    "\n",
    "    def update(i):\n",
    "        fig.clear()\n",
    "        im = plt.scatter(zs[i][:, 0], zs[i][:, 1], c=c, cmap=cm.coolwarm)\n",
    "        im.axes.set_title(\"Epoch: \" + str(epochs[i]))\n",
    "        im.axes.set_xlim(-5, 5)\n",
    "        im.axes.set_ylim(-5, 5)\n",
    "        return im\n",
    "\n",
    "    anim = FuncAnimation(fig, update, frames=range(len(zs)), interval=150)\n",
    "    anim.save(fname, dpi=80, writer='imagemagick')\n",
    "    \n",
    "make_2d_figs_gif(figs, epochs, \"./figs3/manifold.gif\", plt.figure(figsize=(10,10)))\n",
    "make_2d_scatter_gif(latent_distrs, epochs, y_test, \"./figs3/z_distr.gif\", plt.figure(figsize=(10,10)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В следующей части посмотрим как генерировать цифры нужного лейбла, а также как переносить стиль с одной цифры на другую."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Полезные ссылки и литература"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Теоретическая часть основана на статье:  \n",
    "[1] Tutorial on Variational Autoencoders, Carl Doersch, 2016, https://arxiv.org/abs/1606.05908  \n",
    "и фактически является ее кратким изложением\n",
    "\n",
    "Многие картинки взяты из блога Isaac Dykeman:  \n",
    "[2] Isaac Dykeman, http://ijdykeman.github.io/ml/2016/12/21/cvae.html  \n",
    "\n",
    "Подробнее прочитать про расстояние Кульбака-Лейблера на русском можно здесь   \n",
    "[3] http://www.machinelearning.ru/wiki/images/d/d0/BMMO11_6.pdf  \n",
    "\n",
    "Код частично основан на статье *Francois Chollet*:  \n",
    "[4] https://blog.keras.io/building-autoencoders-in-keras.html  \n",
    "\n",
    "Другие интересные ссылки:  \n",
    "http://blog.fastforwardlabs.com/2016/08/12/introducing-variational-autoencoders-in-prose-and.html  \n",
    "http://kvfrans.com/variational-autoencoders-explained/  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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
}