analyse_trained-baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# export THEANO_FLAGS=\"device=gpu0, floatX=float32\" optimizer=None\n",
    "\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import numpy as np\n",
    "import os\n",
    "import socket\n",
    "import argparse\n",
    "import time\n",
    "import datetime\n",
    "import importlib\n",
    "import pprint\n",
    "\n",
    "%matplotlib inline\n",
    "# %matplotlib nbagg\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "\n",
    "import lasagne\n",
    "from lasagne.utils import floatX\n",
    "from lasagne.updates import rmsprop, adam, momentum\n",
    "from lasagne.layers import get_all_params, get_all_layers, get_all_param_values, get_output\n",
    "from lasagne.objectives import squared_error, binary_crossentropy, aggregate\n",
    "\n",
    "from utils.helperFunctions import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# -------- setup options and data ------------------\n",
    "pretrained_model_path = '/path/to/model_checkpoint.p'\n",
    "\n",
    "checkpoint = pickle.load(open(pretrained_model_path, 'rb'))\n",
    "model_values = checkpoint['model_values'] # overwrite the values of model parameters\n",
    "options = checkpoint['options']\n",
    "pprint.PrettyPrinter(indent=4).pprint(options)\n",
    "\n",
    "# Load options\n",
    "np.random.seed(options['seed'])\n",
    "host = socket.gethostname() # get computer hostname\n",
    "start_time = datetime.datetime.now().strftime(\"%y-%m-%d-%H-%M\")\n",
    "\n",
    "model = importlib.import_module(options['model_file'])\n",
    "\n",
    "# Optional: change some options\n",
    "# options['modelOptions']['target_seqlen'] = 10\n",
    "# options['datasetOptions']['num_frames'] = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# ---------- build model and compile ---------------\n",
    "input_batch = T.tensor4() # input image sequences\n",
    "target = T.tensor4() # target image\n",
    "\n",
    "print('Build model...')\n",
    "model = model.Model(**options['modelOptions'])\n",
    "\n",
    "print('Compile ...')\n",
    "net, outputs = model.build_model(input_batch)\n",
    "\n",
    "# compute loss\n",
    "outputs = get_output(outputs)\n",
    "output_frames = outputs\n",
    "\n",
    "# # compute loss\n",
    "# outputs = get_output(outputs + [filters])\n",
    "# output_frames = outputs[:-1]\n",
    "# output_filter = outputs[-1]\n",
    "\n",
    "train_losses = []\n",
    "for i in range(options['modelOptions']['target_seqlen']):\n",
    "    output_frame = output_frames[i]\n",
    "\n",
    "    if options['loss'] == 'squared_error':\n",
    "        frame_loss = squared_error(output_frame, target[:, [i], :, :])\n",
    "    elif options['loss'] == 'binary_crossentropy':\n",
    "        # Clipping to avoid NaN's in binary crossentropy: https://github.com/Lasagne/Lasagne/issues/436\n",
    "        output_frame = T.clip(output_frame, np.finfo(np.float32).eps, 1-np.finfo(np.float32).eps)\n",
    "        frame_loss = binary_crossentropy(output_frame, target[:,[i],:,:])\n",
    "    else:\n",
    "        assert False\n",
    "\n",
    "    train_losses.append(aggregate(frame_loss))\n",
    "\n",
    "train_loss = sum(train_losses) / options['modelOptions']['target_seqlen']\n",
    "\n",
    "# update\n",
    "sh_lr = theano.shared(lasagne.utils.floatX(options['learning_rate'])) # to allow dynamic learning rate\n",
    "\n",
    "layers = get_all_layers(net)\n",
    "all_params = get_all_params(layers, trainable = True)\n",
    "updates = adam(train_loss, all_params, learning_rate=sh_lr)\n",
    "_test = theano.function([input_batch, target], [train_loss] + output_frames, allow_input_downcast=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# ------------ data setup ----------------\n",
    "print('Prepare data...')\n",
    "if options['dataset_file'] == 'datasets.bouncingMnist_original':\n",
    "    options['dataset_file'] = 'datasets.bouncingMnist_originalTest'\n",
    "\n",
    "datasetOptions = options['datasetOptions']\n",
    "dataset = importlib.import_module(options['dataset_file'])\n",
    "datasetOptions['mode'] = 'test'\n",
    "dh = dataset.DataHandler(**datasetOptions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# ------------ training setup ----------------\n",
    "lasagne.layers.set_all_param_values(layers, model_values)\n",
    "history_train = checkpoint['history_train']\n",
    "batch_size = options['batch_size']\n",
    "\n",
    "# ------------ actual training ----------------\n",
    "input_seqlen = options['modelOptions']['input_seqlen']\n",
    "\n",
    "# setup a test batch\n",
    "its = dh.GetDatasetSize() // options['batch_size']\n",
    "# its = 1\n",
    "history_batch = []\n",
    "for i in range(0, its):\n",
    "    ind = np.arange(i*batch_size, (i+1)*batch_size)\n",
    "#     import pdb; pdb.set_trace()\n",
    "    batch = dh.GetBatch(ind) # generate data on the fly\n",
    "    if options['dataset_file'] == 'datasets.stereoCarsColor':\n",
    "        batch_input = batch[..., :input_seqlen].squeeze(axis=4)  # first frames\n",
    "        batch_target = batch[..., input_seqlen:].squeeze(axis=4)  # last frame\n",
    "    else:\n",
    "        batch_input = batch[..., :input_seqlen].transpose(0,4,2,3,1).squeeze(axis=4) # first frames\n",
    "        batch_target = batch[..., input_seqlen:].transpose(0,4,2,3,1).squeeze(axis=4) # last frame\n",
    "    testOutputs = _test(batch_input, batch_target)\n",
    "    loss_test = testOutputs[0]\n",
    "    history_batch.append(loss_test)\n",
    "#     import pdb; pdb.set_trace()\n",
    "    predictions = np.asarray(testOutputs[1:]).transpose(1,2,3,4,0)\n",
    "    print(\"Batch {} of {}\".format(i+1, its))\n",
    "    print(\"Test loss:\\t{:.6f}\".format(np.mean(loss_test)))\n",
    "    \n",
    "    savefigs = False\n",
    "    if savefigs == True:\n",
    "        for case_id in range(batch_size):\n",
    "#             import pdb; pdb.set_trace()\n",
    "            # clear the screen\n",
    "            display.clear_output(wait=True)\n",
    "            \n",
    "            # visualize the prediction\n",
    "            visualize_prediction(batch, fut=predictions, fig=1, case_id=case_id, saveId=(i+1)*100+case_id, savefig=True)\n",
    "\n",
    "            # visualize the flow map\n",
    "            visualize_flowmap(pred_filter, batch, predictions, input_seqlen, options['image_dim'], options['modelOptions']['dynamic_filter_size'][0], case_id=case_id, saveId=(i+1)*100+case_id, savefig=True)\n",
    "#             visualize_flowmapStereo(pred_filter, batch, predictions, input_seqlen, options['image_dim'], options['modelOptions']['dynamic_filter_size'][0], case_id=case_id, saveId=(i+1)*100+case_id, savefig=True)\n",
    "\n",
    "            # animated gif\n",
    "            import matplotlib.animation as animation\n",
    "\n",
    "            fig = plt.figure() # make figure\n",
    "\n",
    "            # make axesimage object\n",
    "            # the vmin and vmax here are very important to get the color map correct\n",
    "            redgreen = np.zeros((3,options['image_dim'],options['image_dim'],options['datasetOptions']['num_frames']))\n",
    "            # import pdb; pdb.set_trace()\n",
    "            redgreen[0,:,:,:input_seqlen] = 1\n",
    "            redgreen[1,:,:,input_seqlen:] = 1\n",
    "\n",
    "            data = batch[case_id]\n",
    "\n",
    "            data2 = np.concatenate((np.zeros((options['batch_size'], 1, options['image_dim'], \n",
    "                                                    options['image_dim'], input_seqlen)), \n",
    "                                          predictions), axis=4)\n",
    "            data2 = np.concatenate((batch[..., :input_seqlen], \n",
    "                                          predictions), axis=4)\n",
    "            data2 = data2[case_id]\n",
    "\n",
    "            plt.subplot(1,3,1)\n",
    "            im0 = plt.imshow(redgreen[..., 0].transpose(1,2,0).squeeze(), cmap=plt.cm.gray, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "            plt.axis('off')\n",
    "            plt.subplot(1,3,2)\n",
    "            im = plt.imshow(data[..., 0].transpose(1,2,0).squeeze(), cmap=plt.cm.gray, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "            plt.axis('off')\n",
    "            plt.subplot(1,3,3)\n",
    "            im2 = plt.imshow(data2[..., 0].transpose(1,2,0).squeeze(), cmap=plt.cm.gray, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "            plt.axis('off')\n",
    "\n",
    "            # function to update figure\n",
    "            def updatefig(j):\n",
    "                # set the data in the axesimage object \n",
    "                im0.set_array(redgreen[..., j].transpose(1,2,0).squeeze())\n",
    "                im.set_array(data[..., j].transpose(1,2,0).squeeze())\n",
    "                im2.set_array(data2[..., j].transpose(1,2,0).squeeze())\n",
    "                return im, im2\n",
    "\n",
    "            # kick off the animation\n",
    "            ani = animation.FuncAnimation(fig, updatefig, frames=range(options['datasetOptions']['num_frames']), interval=200, repeat_delay=1000, blit=True)\n",
    "            ani.save('images/%d' % ((i+1)*100+case_id) + '.gif', writer='imagemagick')\n",
    "            plt.show()\n",
    "# print statistics\n",
    "print(\"  Test loss:\\t{:.6f}\".format(np.mean(history_batch)))\n",
    "print(\"  Parameter count: {}\".format(lasagne.layers.count_params(net)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "savefig = False\n",
    "case_id = 0 # element of the batch to use for visualization\n",
    "\n",
    "# convergence plot\n",
    "plt.figure()\n",
    "plt.plot(range(1,len(history_train)+1), history_train, label=\"loss\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# visualize the prediction\n",
    "visualize_prediction(batch, fut=predictions, fig=1, case_id=case_id, savefig=savefig)\n",
    "    \n",
    "visualize_flowmap(pred_filter, batch, predictions, input_seqlen, options['image_dim'], options['modelOptions']['dynamic_filter_size'][0], case_id, savefig=savefig)\n",
    "# visualize_flowmapStereo(pred_filter, batch, predictions, input_seqlen, options['image_dim'], options['modelOptions']['dynamic_filter_size'][0], case_id, savefig=savefig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.animation as animation\n",
    "\n",
    "fig = plt.figure() # make figure\n",
    "case_id = 1\n",
    "\n",
    "# make axesimage object\n",
    "# the vmin and vmax here are very important to get the color map correct\n",
    "redgreen = np.zeros((3,options['image_dim'],options['image_dim'],options['datasetOptions']['num_frames']))\n",
    "# import pdb; pdb.set_trace()\n",
    "redgreen[0,:,:,:input_seqlen] = 1\n",
    "redgreen[1,:,:,input_seqlen:] = 1\n",
    "\n",
    "data = batch[case_id]\n",
    "\n",
    "data2 = np.concatenate((np.zeros((options['batch_size'], 1, options['image_dim'], \n",
    "                                        options['image_dim'], input_seqlen)), \n",
    "                              predictions), axis=4)\n",
    "data2 = np.concatenate((batch[..., :input_seqlen], \n",
    "                              predictions), axis=4)\n",
    "data2 = data2[case_id]\n",
    "\n",
    "plt.subplot(1,3,1)\n",
    "im0 = plt.imshow(redgreen[..., 0].transpose(1,2,0).squeeze(), cmap=plt.cm.gray, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "plt.axis('off')\n",
    "plt.subplot(1,3,2)\n",
    "im = plt.imshow(data[..., 0].transpose(1,2,0).squeeze(), cmap=plt.cm.gray, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "plt.axis('off')\n",
    "plt.subplot(1,3,3)\n",
    "im2 = plt.imshow(data2[..., 0].transpose(1,2,0).squeeze(), cmap=plt.cm.gray, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "plt.axis('off')\n",
    "\n",
    "# function to update figure\n",
    "def updatefig(j):\n",
    "    # set the data in the axesimage object \n",
    "    im0.set_array(redgreen[..., j].transpose(1,2,0).squeeze())\n",
    "    im.set_array(data[..., j].transpose(1,2,0).squeeze())\n",
    "    im2.set_array(data2[..., j].transpose(1,2,0).squeeze())\n",
    "    return im, im2\n",
    "\n",
    "# kick off the animation\n",
    "ani = animation.FuncAnimation(fig, updatefig, frames=range(options['datasetOptions']['num_frames']), interval=100, repeat_delay=1000, blit=True)\n",
    "ani.save(options['name'] + '.gif', writer='imagemagick')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}