experiments/many_well/fab_many_well.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/lollcat/fab-torch/blob/master/experiments/many_well/fab_many_well.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XpTfsy3InUhf"
   },
   "source": [
    "# Flow Annealed Importance Sampling Bootstrap: Many Well Problem\n",
    "In this notebook we will compare training a flow using FAB with a prioritised buffer, to training a flow by reverse KL divergence minimisation. We will train the models relatively briefly to get indications of how well each method works in a relatively small amount of time, however better results may be obtained by simply increasing the training time. In this notebook we train a flow on a 6 dimensional version of the Many Well problem. The problem difficulty may be increased by increasing the dimension of the Many Well problem.\n",
    "\n",
    "GPU is not required for this notebook. Each experiment runs on my laptop (CPU only) in under 2 minuates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pIuF7gAmLbpI"
   },
   "source": [
    "# Setup Repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U21CxXjHRcsF"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/lollcat/fab-torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tmsNqL0lRwGa"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir(\"fab-torch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d7eC5CsoRs_d"
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys\n",
    "# sys.path.insert(0, \"../\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1Ajs-kTgLeWU"
   },
   "source": [
    "# Let's go!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TeFLDH5mLhv9"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lzkmrn81LajP"
   },
   "outputs": [],
   "source": [
    "import normflows as nf\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from fab import FABModel, HamiltonianMonteCarlo, Metropolis\n",
    "from fab.utils.logging import ListLogger\n",
    "from fab.utils.plotting import plot_history, plot_contours, plot_marginal_pair\n",
    "from fab.target_distributions.many_well import ManyWellEnergy\n",
    "from fab.utils.prioritised_replay_buffer import PrioritisedReplayBuffer\n",
    "from fab import Trainer, PrioritisedBufferTrainer\n",
    "from fab.utils.plotting import plot_contours, plot_marginal_pair\n",
    "\n",
    "\n",
    "from experiments.make_flow import make_wrapped_normflow_realnvp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OHtMPbFlMKvd"
   },
   "source": [
    "## Setup Target distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kkr2CqqDMRBn"
   },
   "outputs": [],
   "source": [
    "dim = 6 # Can increase in to higher values that are multiples of two.\n",
    "seed = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qM9PUDE3MMoA"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(0)  # seed of 0 for GMM problem\n",
    "target = ManyWellEnergy(dim, a=-0.5, b=-6, use_gpu=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oxMmszREPEY1"
   },
   "outputs": [],
   "source": [
    "# plot the contours for the marginal distribution of the first 2D of target (i.e. the Double Well Problem). \n",
    "target.to(\"cpu\")\n",
    "fig, ax = plt.subplots()\n",
    "plotting_bounds = (-3, 3)\n",
    "plot_contours(target.log_prob_2D, bounds=plotting_bounds, n_contour_levels=40, ax=ax, grid_width_n_points=100)\n",
    "if torch.cuda.is_available():\n",
    "    target.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bJmoBOJ8REZO"
   },
   "source": [
    "## Create FAB model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xp58k3FMQ3Qf"
   },
   "outputs": [],
   "source": [
    "# hyper-parameters\n",
    "\n",
    "# Flow\n",
    "n_flow_layers = 10\n",
    "layer_nodes_per_dim = 40\n",
    "lr = 2e-4\n",
    "max_gradient_norm = 100.0\n",
    "batch_size = 128\n",
    "n_iterations = 500\n",
    "n_eval = 10\n",
    "eval_batch_size = batch_size * 10\n",
    "n_plots = 10 # number of plots shows throughout tranining\n",
    "use_64_bit = True\n",
    "alpha = 2.0\n",
    "\n",
    "# AIS\n",
    "transition_operator_type = \"hmc\"\n",
    "n_intermediate_distributions = 1\n",
    "\n",
    "# buffer config\n",
    "n_batches_buffer_sampling = 4\n",
    "maximum_buffer_length = batch_size* 100\n",
    "min_buffer_length = batch_size * 10\n",
    "\n",
    "# target p^\\alpha q^{a-\\alpha} as target for AIS. \n",
    "min_is_target = True\n",
    "p_target = not min_is_target # Whether to use p as the target. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3P5c29Rayd2B"
   },
   "outputs": [],
   "source": [
    "if use_64_bit:\n",
    "    torch.set_default_dtype(torch.float64)\n",
    "    target = target.double()\n",
    "    print(f\"running with 64 bit\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MRJx0FhTRKIF"
   },
   "source": [
    "### Setup flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ptJrkMn5Qz2F"
   },
   "outputs": [],
   "source": [
    "flow = make_wrapped_normflow_realnvp(dim, n_flow_layers=n_flow_layers, \n",
    "                                 layer_nodes_per_dim=layer_nodes_per_dim,\n",
    "                                act_norm = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bNhmNi0zRMT2"
   },
   "source": [
    "### Setup Transition operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aAfUX8rgQ9XG"
   },
   "outputs": [],
   "source": [
    "if transition_operator_type == \"hmc\":\n",
    "    # very lightweight HMC.\n",
    "    transition_operator = HamiltonianMonteCarlo(\n",
    "            n_ais_intermediate_distributions=n_intermediate_distributions,\n",
    "            dim=dim,\n",
    "            base_log_prob=flow.log_prob,\n",
    "            target_log_prob=target.log_prob,\n",
    "            alpha=alpha,\n",
    "            p_target=p_target,\n",
    "            n_outer=1,\n",
    "            epsilon=1.0, \n",
    "            L=5)\n",
    "elif transition_operator_type == \"metropolis\":\n",
    "    transition_operator = Metropolis(            \n",
    "        n_ais_intermediate_distributions=n_intermediate_distributions,\n",
    "        dim=dim,\n",
    "        base_log_prob=flow.log_prob,\n",
    "        target_log_prob=target.log_prob,\n",
    "        alpha=alpha,\n",
    "        p_target=p_target,\n",
    "        n_updates=1,\n",
    "        adjust_step_size=False,\n",
    "        max_step_size=metropolis_step_size, # the same for all metropolis steps \n",
    "        min_step_size=metropolis_step_size,\n",
    "        eval_mode=False,\n",
    "                                  )\n",
    "else:\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oIQthDLyLkus"
   },
   "source": [
    "### Setup FAB model with prioritised replay buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wFrJvytJcAm2"
   },
   "outputs": [],
   "source": [
    "# use GPU if available\n",
    "if torch.cuda.is_available():\n",
    "    flow.cuda()\n",
    "    transition_operator.cuda()\n",
    "    print(f\"Running with GPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SgXAZZpCSAiK"
   },
   "outputs": [],
   "source": [
    "fab_model = FABModel(flow=flow,\n",
    "                     target_distribution=target,\n",
    "                     n_intermediate_distributions=n_intermediate_distributions,\n",
    "                     transition_operator=transition_operator,\n",
    "                     alpha=alpha\n",
    "                    )\n",
    "optimizer = torch.optim.Adam(flow.parameters(), lr=lr)\n",
    "logger = ListLogger(save=False) # save training history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5onsUCTJbE-l"
   },
   "outputs": [],
   "source": [
    "# Setup buffer.\n",
    "def initial_sampler():\n",
    "  # fill replay buffer using initialised model and AIS.\n",
    "    point, log_w = fab_model.annealed_importance_sampler.sample_and_log_weights(\n",
    "            batch_size, logging=False)\n",
    "    return point.x, log_w, point.log_q\n",
    "buffer = PrioritisedReplayBuffer(dim=dim, max_length=maximum_buffer_length,\n",
    "                      min_sample_length=min_buffer_length,\n",
    "                      initial_sampler=initial_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B2QJeaj5Ll4o"
   },
   "outputs": [],
   "source": [
    "def plot(fab_model, n_samples = batch_size, dim=dim):\n",
    "    n_rows = dim // 2\n",
    "    fig, axs = plt.subplots(dim // 2, 2,  sharex=True, sharey=True, figsize=(10, n_rows*3))\n",
    "\n",
    "\n",
    "    samples_flow = fab_model.flow.sample((n_samples,))\n",
    "    samples_ais = fab_model.annealed_importance_sampler.sample_and_log_weights(n_samples,\n",
    "                                                                               logging=False)[0].x\n",
    "\n",
    "    for i in range(n_rows):\n",
    "        plot_contours(target.log_prob_2D, bounds=plotting_bounds, ax=axs[i, 0], n_contour_levels=40)\n",
    "        plot_contours(target.log_prob_2D, bounds=plotting_bounds, ax=axs[i, 1], n_contour_levels=40)\n",
    "\n",
    "        # plot flow samples\n",
    "        plot_marginal_pair(samples_flow, ax=axs[i, 0], bounds=plotting_bounds, marginal_dims=(i*2,i*2+1))\n",
    "        axs[i, 0].set_xlabel(f\"dim {i*2}\")\n",
    "        axs[i, 0].set_ylabel(f\"dim {i*2 + 1}\")\n",
    "\n",
    "\n",
    "\n",
    "        # plot ais samples\n",
    "        plot_marginal_pair(samples_ais, ax=axs[i, 1], bounds=plotting_bounds, marginal_dims=(i*2,i*2+1))\n",
    "        axs[i, 1].set_xlabel(f\"dim {i*2}\")\n",
    "        axs[i, 1].set_ylabel(f\"dim {i*2+1}\")\n",
    "        plt.tight_layout()\n",
    "    axs[0, 1].set_title(\"ais samples\")  \n",
    "    axs[0, 0].set_title(\"flow samples\")\n",
    "    plt.show()\n",
    "    return [fig]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ecjfGOS2bWEq"
   },
   "outputs": [],
   "source": [
    "plot(fab_model) # Visualise model during initialisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zfjeD4Udb275"
   },
   "outputs": [],
   "source": [
    "# Setup trainer.\n",
    "trainer = PrioritisedBufferTrainer(model=fab_model, optimizer=optimizer, \n",
    "                                   logger=logger, plot=plot,\n",
    "                        buffer=buffer, n_batches_buffer_sampling=n_batches_buffer_sampling,\n",
    "                     max_gradient_norm=max_gradient_norm, alpha=alpha)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ailsWaOwdF5V"
   },
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5Fyi-oe8cbXz"
   },
   "outputs": [],
   "source": [
    "# Now run!\n",
    "trainer.run(n_iterations=n_iterations, batch_size=batch_size, n_plot=n_plots, \\\n",
    "            n_eval=n_eval, eval_batch_size=eval_batch_size, save=False) # note that the progress bar during training prints ESS w.r.t p^2/q. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "puWE84Itm77J"
   },
   "source": [
    "In the below plot of samples from the flow vs the target contours, and with the test set log prob throughout training, we see that the flow covers the target distribution quite well. It may be trained further to obtain even better results. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gGsDYvqY23TK"
   },
   "outputs": [],
   "source": [
    "# \"_eval\" means metrics calculated with eval_batch_size, _p_target means metrics calculated with AIS targetting p, p2overq_target means calculated with AIS targeting p^2/q. \n",
    "# For example 'eval_ess_flow_p2overq_target' is the effective sample size of the flow w.r.t the target distribution p^2/q when sampling from AIS with p^2/q as the target.\n",
    "logger.history.keys() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lce_nsLdkKSB"
   },
   "outputs": [],
   "source": [
    "# Test set probability using samples from the target distribution.\n",
    "eval_iters = np.linspace(0, n_iterations, n_eval)\n",
    "plt.plot(eval_iters, logger.history['test_set_exact_mean_log_prob_p_target'])\n",
    "plt.ylabel(\"mean test set log prob\")\n",
    "plt.xlabel(\"training iteration\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w31pgTdNYCOh"
   },
   "outputs": [],
   "source": [
    "# Effective sample size\n",
    "eval_iters = np.linspace(0, n_iterations, n_eval)\n",
    "plt.plot(eval_iters, logger.history['eval_ess_flow_p_target'])\n",
    "plt.ylabel(\"Effective Sample Size\")\n",
    "plt.xlabel(\"training iteration\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lxe2AYrLYSaj"
   },
   "outputs": [],
   "source": [
    "# Probability of test set containing a point on each mode\n",
    "eval_iters = np.linspace(0, n_iterations, n_eval)\n",
    "plt.plot(eval_iters, logger.history['test_set_modes_mean_log_prob_p_target'])\n",
    "plt.ylabel(\"Average log prob of modes test set\")\n",
    "plt.xlabel(\"training iteration\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BLiWtIEw3VqK"
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 1, figsize=(5, 5))\n",
    "target.to(\"cpu\")\n",
    "plot_contours(target.log_prob_2D, bounds=plotting_bounds, ax=axs, n_contour_levels=40, grid_width_n_points=200)\n",
    "if torch.cuda.is_available():\n",
    "    target.to(\"cuda\")\n",
    "\n",
    "n_samples = 1000\n",
    "samples_flow = fab_model.flow.sample((n_samples,)).detach()\n",
    "plot_marginal_pair(samples_flow, ax=axs, bounds=plotting_bounds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8ajP0mo4hjjr"
   },
   "source": [
    "# Training a flow by reverse KL divergence minimisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aXBZFZT11dC_"
   },
   "outputs": [],
   "source": [
    "loss_type = \"flow_reverse_kl\" # can set to \"target_foward_kl\" for training by maximum likelihood of samples from the Many Well target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LL_MnSZrhnT0"
   },
   "outputs": [],
   "source": [
    "# Create flow using the same architecture.\n",
    "flow = make_wrapped_normflow_realnvp(dim, n_flow_layers=n_flow_layers, \n",
    "                                 layer_nodes_per_dim=layer_nodes_per_dim,\n",
    "                                act_norm = False)\n",
    "optimizer = torch.optim.Adam(flow.parameters(), lr=lr)\n",
    "logger = ListLogger(save=False) # save training history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EtOPcNhqrajL"
   },
   "outputs": [],
   "source": [
    "# use GPU if available\n",
    "if torch.cuda.is_available():\n",
    "    flow.cuda()\n",
    "    transition_operator.cuda()\n",
    "    print(f\"Running with GPU\")\n",
    "    target.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8_OlZmP9ugI1"
   },
   "outputs": [],
   "source": [
    "n_iterations = int(4*(n_iterations)) # Training the flow by KL minimisation is cheaper per iteration, so we run it for more iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7nZPakrl1g2-"
   },
   "outputs": [],
   "source": [
    "reverse_kld_model = FABModel(flow=flow,\n",
    "                     target_distribution=target,\n",
    "                     n_intermediate_distributions=n_intermediate_distributions,\n",
    "                     transition_operator=transition_operator,\n",
    "                     loss_type=loss_type,\n",
    "                     alpha=None\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sDQ29z5t2zSa"
   },
   "outputs": [],
   "source": [
    "def plot_flow_reverse_kld(fab_model, n_samples = batch_size, dim=dim):\n",
    "    n_rows = dim // 2\n",
    "    fig, axs = plt.subplots(dim // 2, 1,  sharex=True, sharey=True, figsize=(5, n_rows*3))\n",
    "\n",
    "    \n",
    "    samples_flow = fab_model.flow.sample((n_samples,))\n",
    "\n",
    "    for i in range(n_rows):\n",
    "      plot_contours(target.log_prob_2D, bounds=plotting_bounds, ax=axs[i], n_contour_levels=40)\n",
    "\n",
    "      # plot flow samples\n",
    "      plot_marginal_pair(samples_flow, ax=axs[i], bounds=plotting_bounds, marginal_dims=(i*2,i*2+1))\n",
    "      axs[i].set_xlabel(f\"dim {i*2}\")\n",
    "      axs[i].set_ylabel(f\"dim {i*2 + 1}\")\n",
    "      plt.tight_layout()\n",
    "    plt.show()\n",
    "    return [fig]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QYAWr5ichwLS"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(model=reverse_kld_model, optimizer=optimizer, logger=logger, plot=plot_flow_reverse_kld, max_gradient_norm=max_gradient_norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "23JzOSosjUn3"
   },
   "outputs": [],
   "source": [
    "# Now run!\n",
    "trainer.run(n_iterations=n_iterations, batch_size=batch_size, n_plot=n_plots, \\\n",
    "            n_eval=n_eval, eval_batch_size=eval_batch_size, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GAUDI5lNweY_"
   },
   "source": [
    "We evaluate the flow on samples from the target distribution, we see that because the flow trained by kl divergence minimisation is missing modes, the flow places NaN log prob on samples from the target. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kTDrcaa5wK9g"
   },
   "outputs": [],
   "source": [
    "logger.history[\"test_set_exact_mean_log_prob\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cyLoMV8VvOSE"
   },
   "outputs": [],
   "source": [
    "eval_iters = np.linspace(0, n_iterations, n_eval)\n",
    "plt.plot(eval_iters, logger.history[\"test_set_exact_mean_log_prob\"])\n",
    "plt.ylabel(\"mean test set log prob\")\n",
    "plt.xlabel(\"eval iteration\")\n",
    "plt.yscale(\"symlog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8ow35UmcU4_X"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "provenance": []
  },
  "gpuClass": "standard",
  "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}