assignment07_starter.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d11f3eac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1st part: implement coupling layer\n",
    "# - split input\n",
    "# - implement y1 = x1, y2 = x2 + m(y1)   (forward)\n",
    "# - implement backward\n",
    "\n",
    "\n",
    "# one way to split could be\n",
    "x1, x2 = tf.split(x, 2, axis=-1)\n",
    "y1, y2 = whatever_the_coupling_layer_does(x1, x2)\n",
    "y =  tf.concat([y1, y2], axis=-1)\n",
    "\n",
    "# this is not so good, since it splits the data right through the middle.\n",
    "# if the data has spatial relations, like images, this means the model\n",
    "# conditions, say, the upper half of the image on the lower half, and vice versa.\n",
    "# these long-range dependencies are hard to model.\n",
    "\n",
    "# another option could be an even-odd split:\n",
    "def split_even_odd(inp):\n",
    "    even_inds = tf.range(0, inp.shape[1], 2)\n",
    "    odd_inds = tf.range(1, inp.shape[1], 2)\n",
    "\n",
    "    even = tf.gather(inp, even_inds, axis=1)\n",
    "    odd = tf.gather(inp, odd_inds, axis=1)\n",
    "\n",
    "    # process even and odd parts separately...\n",
    "    ...\n",
    "    \n",
    "    # combine\n",
    "    together = tf.stack([even_output, odd_output], axis=-1)\n",
    "    return tf.reshape(together, tf.shape(inp))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1019f718",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow_probability as tfp\n",
    "tfd = tfp.distributions\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d11fc6a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2nd part: NICE model\n",
    "# - stack a bunch of coupling layers\n",
    "# - switch x1, x2 for each layer\n",
    "# - forward: apply coupling layers\n",
    "# - backward: apply backward coupling layers (in reverse)\n",
    "\n",
    "# - rescaling:\n",
    "# - create d-dimensional vector (model weight)\n",
    "# - forward: multiply x * exp(vector) after applying coupling layers\n",
    "# - backwards: x / vector (or x * tf.exp(-vector)) BEFORE applying reverse coupling layers\n",
    "\n",
    "\n",
    "# - training: map x -> h using NICE model\n",
    "# - compute log_p_simple(h)\n",
    "# - add log determinant of jacobian: simply sum of scaling values\n",
    "# use -log_likelihood as loss\n",
    "\n",
    "# a note on simple distributions.\n",
    "# you should use tfd = tfp.distributions...\n",
    "input_dim = 12  # example\n",
    "batch_size = 8\n",
    "dummy_data = tf.random.normal((batch_size, input_dim))\n",
    "\n",
    "simple_distribution = tfd.Normal(loc=tf.zeros((input_dim,)), scale=tf.ones((input_dim,)))\n",
    "log_p_simple = simple_distribution.log_prob(dummy_data)\n",
    "# ... this will return a batch x dim matrix. then SUM over axis 1 (data dimension)! average over batch axis!\n",
    "log_p_simple = tf.reduce_sum(log_p_simple, axis=1)\n",
    "print(log_p_simple)\n",
    "\n",
    "# alternatively: use multivariate distribution.\n",
    "# this returns one prob per entry.\n",
    "simple_distribution = tfd.MultivariateNormalDiag(loc=tf.zeros((input_dim,)))\n",
    "print(simple_distribution.log_prob(dummy_data))\n",
    "\n",
    "print(\"Same results!\")\n",
    "\n",
    "# just don't mix it up!!! I had a bug where I was using multivariate, but summing over the last dimension.\n",
    "# this resulted in summing my loss over the batch axis instead of averaging.\n",
    "# that's bad, because your effective learning rate is MUCH higher than expected.\n",
    "\n",
    "# the paper proses using tfd.Logistic instead. that does not have a multivariate version AFAIK,\n",
    "# so you will have to use the first option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8fb9596",
   "metadata": {},
   "outputs": [],
   "source": [
    "# finally, here is a sample toy dataset.\n",
    "# it's a 2D parabola kinda thing.\n",
    "# you can use this to test your flow models.\n",
    "# even a simple model like NICE should be able to fit this!\n",
    "n_samples = 2048\n",
    "x2_dist = tfd.Normal(loc=0., scale=0.5)\n",
    "x2_samples = x2_dist.sample(n_samples)\n",
    "x1 = tfd.Normal(loc=1. * tf.square(x2_samples),\n",
    "                scale=0.1*tf.ones(n_samples, dtype=tf.float32))\n",
    "x1_samples = x1.sample()\n",
    "x_samples = tf.stack([x1_samples, x2_samples], axis=1)\n",
    "\n",
    "as_np = x_samples.numpy()\n",
    "plt.scatter(as_np[:, 0], as_np[:, 1])\n",
    "a = plt.gca()\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# x_samples it the dataset!\n",
    "# NOTE that I say in the assignment that LayerNorm > BatchNorm for these models apparently.\n",
    "# but for these simple models with very low data dimensionality, it appears layernorm causes issues sometimes.\n",
    "# so maybe leave normalization out completely."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}