colab_example_gaussian_mixture.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dcZL4LHRnZ6v",
    "outputId": "a06097fc-4c4c-4335-a10a-9fb44592e106"
   },
   "source": [
    "# Example: Gaussian Mixture\n",
    "\n",
    "2020-11-29 first uploaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/jaekookang/invertible_neural_networks.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OidApFl0n3ll",
    "outputId": "f7978b66-c400-47ac-a3d9-c0528f54fc07"
   },
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "import numpy as np\n",
    "import itertools\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import tensorflow as tf\n",
    "print(tf.__version__)\n",
    "!python --version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OidApFl0n3ll",
    "outputId": "f7978b66-c400-47ac-a3d9-c0528f54fc07"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, 'invertible_neural_networks')\n",
    "from flow import *\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bL7OG5h4ng3J"
   },
   "outputs": [],
   "source": [
    "n_means = 8\n",
    "radius = 14\n",
    "sd = 1\n",
    "labels = ['red','gold','green','chocolate','blue','magenta','pink','purple']\n",
    "# labels = ['red','red','red','red','blue','blue','green','purple']\n",
    "assert len(labels) == n_means\n",
    "\n",
    "x_dim = 2\n",
    "y_dim = len(list(set(labels)))\n",
    "z_dim = 2\n",
    "tot_dim = y_dim + z_dim\n",
    "pad_dim = tot_dim - x_dim\n",
    "n_sample = 200\n",
    "n_data = n_sample * n_means\n",
    "n_couple_layer = 3\n",
    "n_hid_layer = 3\n",
    "n_hid_dim = 512\n",
    "\n",
    "n_batch = 200\n",
    "n_epoch = 1000\n",
    "n_display = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 337
    },
    "id": "yHNukZp0n9Z6",
    "outputId": "2799307c-989c-4cc3-d8dd-d8069f28058b"
   },
   "outputs": [],
   "source": [
    "# Make data\n",
    "X_raw = np.zeros((n_means, n_sample, x_dim), dtype='float32')\n",
    "for i in range(n_means):\n",
    "    th = 2*np.pi / n_means * (i+1)\n",
    "    mean = [radius*np.cos(th), radius*np.sin(th)]\n",
    "    X_raw[i, :, :] = np.random.multivariate_normal(mean, np.identity(x_dim)*sd, size=n_sample)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5,5), facecolor='white')\n",
    "for i in range(n_means):\n",
    "    ax.scatter(X_raw[i,:,0], X_raw[i,:,1], s=1)\n",
    "print(X_raw.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 320
    },
    "id": "rG9TmtYw4Bee",
    "outputId": "fabd6cd8-1a1c-4dcc-de92-4f61a9645cab"
   },
   "outputs": [],
   "source": [
    "# Preprocess\n",
    "uq_labels = list(set(labels))\n",
    "idx2lab = {i:lab for i, lab in enumerate(uq_labels)}\n",
    "lab2idx = {idx2lab[key]:i for i, key in enumerate(idx2lab.keys())}\n",
    "\n",
    "X = X_raw.reshape((-1, x_dim))\n",
    "X = StandardScaler().fit_transform(X)\n",
    "y = [[lab2idx[lab]]*n_sample for lab in labels]\n",
    "y = list(itertools.chain.from_iterable(y)) # flatten\n",
    "y_onehot = np.eye(len(uq_labels))[y].astype('int')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5,5), facecolor='white')\n",
    "for i, color in zip(idx2lab.keys(), lab2idx.keys()):\n",
    "    idx = [True if j==i else False for j in y]\n",
    "    ax.scatter(X[idx,0], X[idx,1], s=1, c=color)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k-rNSDc5pyjz"
   },
   "outputs": [],
   "source": [
    "# Pad data\n",
    "pad_x = np.zeros((X.shape[0], pad_dim))\n",
    "x_data = np.concatenate([X, pad_x], axis=-1).astype('float32')\n",
    "z = np.random.multivariate_normal([0.]*x_dim, np.eye(x_dim), X.shape[0])\n",
    "y_data = np.concatenate([z, y_onehot], axis=-1).astype('float32')\n",
    "\n",
    "# Make dataset generator\n",
    "x_data = tf.data.Dataset.from_tensor_slices(x_data)\n",
    "y_data = tf.data.Dataset.from_tensor_slices(y_data)\n",
    "dataset = (tf.data.Dataset.zip((x_data, y_data))\n",
    "           .shuffle(buffer_size=X.shape[0])\n",
    "           .batch(n_batch, drop_remainder=True)\n",
    "           .repeat())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pEAw90SZng-B",
    "outputId": "107f120e-4905-4093-bca8-bffe778b7068"
   },
   "outputs": [],
   "source": [
    "model = NVP(tot_dim, n_couple_layer, n_hid_layer, n_hid_dim, name='NVP')\n",
    "x = tfk.Input((tot_dim,))\n",
    "model(x);\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FKI6S4FIza5w"
   },
   "outputs": [],
   "source": [
    "class Trainer(tfk.Model):\n",
    "    def __init__(self, model, x_dim, y_dim, z_dim, tot_dim, \n",
    "                 n_couple_layer, n_hid_layer, n_hid_dim, shuffle_type='reverse'):\n",
    "        super(Trainer, self).__init__()\n",
    "        self.model = model\n",
    "        self.x_dim = x_dim\n",
    "        self.y_dim = y_dim\n",
    "        self.z_dim = z_dim\n",
    "        self.tot_dim = tot_dim\n",
    "        self.x_pad_dim = tot_dim - x_dim\n",
    "        self.y_pad_dim = tot_dim - (y_dim + z_dim)\n",
    "        self.n_couple_layer = n_couple_layer\n",
    "        self.n_hid_layer = n_hid_layer\n",
    "        self.n_hid_dim = n_hid_dim\n",
    "        self.shuffle_type = shuffle_type\n",
    "\n",
    "        self.w1 = 5.\n",
    "        self.w2 = 1.\n",
    "        self.w3 = 10.\n",
    "        self.loss_factor = 1.\n",
    "        self.loss_fit = MSE\n",
    "        self.loss_latent = MMD_multiscale\n",
    "\n",
    "    def train_step(self, data):\n",
    "        x_data, y_data = data\n",
    "        x = x_data[:, :self.x_dim]\n",
    "        y = y_data[:, -self.y_dim:]\n",
    "        z = y_data[:, :self.z_dim]\n",
    "        y_short = tf.concat([z, y], axis=-1)\n",
    "\n",
    "        # Forward loss\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_out = self.model(x_data)    \n",
    "            pred_loss = self.w1 * self.loss_fit(y_data[:,self.z_dim:], y_out[:,self.z_dim:]) # [zeros, y] <=> [zeros, yhat]\n",
    "            output_block_grad = tf.concat([y_out[:,:self.z_dim], y_out[:, -self.y_dim:]], axis=-1) # take out [z, y] only (not zeros)\n",
    "            latent_loss = self.w2 * self.loss_latent(y_short, output_block_grad) # [z, y] <=> [zhat, yhat]\n",
    "            forward_loss = pred_loss + latent_loss\n",
    "        grads_forward = tape.gradient(forward_loss, self.model.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(grads_forward, self.model.trainable_weights))\n",
    "\n",
    "        # Backward loss\n",
    "        with tf.GradientTape() as tape:\n",
    "            x_rev = self.model.inverse(y_data)\n",
    "            rev_loss = self.w3 * self.loss_factor * self.loss_fit(x_rev, x_data)\n",
    "        grads_backward = tape.gradient(rev_loss, self.model.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(grads_backward, self.model.trainable_weights)) \n",
    "\n",
    "        total_loss = forward_loss + latent_loss + rev_loss\n",
    "        return {'total_loss': total_loss,\n",
    "                'forward_loss': forward_loss,\n",
    "                'latent_loss': latent_loss,\n",
    "                'rev_loss': rev_loss}\n",
    "\n",
    "    def test_step(self, data):\n",
    "        x_data, y_data = data\n",
    "        return NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nUkERU5_zbi4"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(model, x_dim, y_dim, z_dim, tot_dim, n_couple_layer, n_hid_layer, n_hid_dim)\n",
    "trainer.compile(optimizer='Adam')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1K5KckgzzbnJ",
    "outputId": "0abf86fe-cb7b-4d31-b63d-fece392dd25b"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "LossFactor = UpdateLossFactor(n_epoch)\n",
    "logger = NBatchLogger(n_display, n_epoch)\n",
    "hist = trainer.fit(dataset,\n",
    "                   batch_size=n_batch,\n",
    "                   epochs=n_epoch,\n",
    "                   steps_per_epoch=n_data//n_batch, \n",
    "                   callbacks=[logger, LossFactor], \n",
    "                   verbose=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 337
    },
    "id": "rcjLDVopnhLh",
    "outputId": "133a1844-40d7-4fc8-f211-04efa4148b0a"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, facecolor='white', figsize=(8,5))\n",
    "ax.plot(hist.history['total_loss'], 'k.-', label='total_loss')\n",
    "ax.plot(hist.history['forward_loss'], 'b.-', label='forward_loss')\n",
    "ax.plot(hist.history['latent_loss'], 'g.-', label='latent_loss')\n",
    "ax.plot(hist.history['rev_loss'], 'r.-', label='inverse_loss')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 356
    },
    "id": "uII2sXbwnhPL",
    "outputId": "00062611-b08c-4182-f06e-203cdbe45c27"
   },
   "outputs": [],
   "source": [
    "z = np.random.multivariate_normal([1.]*z_dim, np.eye(z_dim), y_onehot.shape[0])\n",
    "y = np.concatenate([z, y_onehot], axis=-1).astype('float32')\n",
    "x_pred = model.inverse(y).numpy()\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10,5), facecolor='white', sharex=True, sharey=True)\n",
    "for i, color in zip(idx2lab.keys(), lab2idx.keys()):\n",
    "    idx = [True if j==i else False for j in y_onehot.argmax(axis=-1)]\n",
    "    ax1.scatter(X[idx,0], X[idx,1], s=1, c=color)\n",
    "    ax2.scatter(x_pred[idx,0], x_pred[idx,1], s=1, c=color)\n",
    "    ax2.set_xlim([-2, 2])\n",
    "    ax2.set_ylim([-2, 2])\n",
    "plt.suptitle('Original (left)                    Prediction (right)', fontsize=20);"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "TMP_inn_multigauss.ipynb",
   "provenance": []
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}