model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, subprocess\n",
    "import tensorflow as tf\n",
    "\n",
    "# config for Gatsby cores\n",
    "if 'gatsby' in str(subprocess.check_output('hostname')):\n",
    "\n",
    "    # specify target gpu device\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = '1' # either '0' or '1' to utilize Titan X GPUs\n",
    "\n",
    "    # allow dynamic GPU allocation\n",
    "    config = tf.ConfigProto()\n",
    "    config.gpu_options.allow_growth = True\n",
    "    session = tf.Session(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# set a seed to control all randomness\n",
    "from tensorflow import set_random_seed\n",
    "from numpy.random import seed\n",
    "\n",
    "set_random_seed(1)\n",
    "seed(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Examine Inputs\n",
    "\n",
    "Let's load some local data and examine its contents. We'll play a small animation that contains the first few temporal frames of the input data below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# X.shape = n_body_parts, n_time_intervals, n_dimensions (3)\n",
    "# each dimension's features are centered 0:1\n",
    "X = np.load('data/npy/carrie-10-mins.npy')\n",
    "\n",
    "# store the shapes of X's values\n",
    "n_vertices, n_time, n_dims = X.shape\n",
    "\n",
    "# load data from pose estimation output\n",
    "#X = np.swapaxes(np.load('data/npy/X.npy'), 0, 1)\n",
    "\n",
    "# labels[i] is a label for the ith body_part\n",
    "labels = ['ARIEL.position', 'C7.position', 'CLAV.position', 'LANK.position', 'LBHD.position', 'LBSH.position', 'LBWT.position', 'LELB.position', 'LFHD.position', 'LFRM.position', 'LFSH.position', 'LFWT.position', 'LHEL.position', 'LIEL.position', 'LIHAND.position', 'LIWR.position', 'LKNE.position', 'LKNI.position', 'LMT1.position', 'LMT5.position', 'LOHAND.position', 'LOWR.position', 'LSHN.position', 'LTHI.position', 'LTOE.position', 'LUPA.position', 'LabelingHips.position', 'MBWT.position', 'MFWT.position', 'RANK.position', 'RBHD.position', 'RBSH.position', 'RBWT.position', 'RELB.position', 'RFHD.position', 'RFRM.position', 'RFSH.position', 'RFWT.position', 'RHEL.position', 'RIEL.position', 'RIHAND.position', 'RIWR.position', 'RKNE.position', 'RKNI.position', 'RMT1.position', 'RMT5.position', 'ROHAND.position', 'ROWR.position', 'RSHN.position', 'RTHI.position', 'RTOE.position', 'RUPA.position', 'STRN.position', 'SolvingHips.position', 'T10.position']\n",
    "\n",
    "# # test that the data is properly centered 0:1 on each dimensional axis\n",
    "# if True:\n",
    "#   for i in range(int(X.shape[2])):\n",
    "#     assert np.min(X[:,:,i]) == 0.0\n",
    "#     assert np.max(X[:,:,i]) == 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import floor\n",
    "\n",
    "# define functions to flatten and unflatten data\n",
    "\n",
    "def flatten(df, run_tests=True):\n",
    "  '''\n",
    "  df is a numpy array with the following three axes:\n",
    "    df.shape[0] = the index of a vertex\n",
    "    df.shape[1] = the index of a time stamp\n",
    "    df.shape[2] = the index of a dimension (x, y, z)\n",
    "  \n",
    "  So df[1][0][2] is the value for the 1st vertex (0-based) at time 0 in dimension 2 (z).\n",
    "  \n",
    "  To flatten this dataframe will mean to push the data into shape:\n",
    "    flattened.shape[0] = time index\n",
    "    flattened.shape[1] = [vertex_index*3] + dimension_vertex\n",
    "    \n",
    "  So flattened[1][3] will be the 3rd dimension of the 1st index (0-based) at time 1. \n",
    "  '''\n",
    "  if run_tests:\n",
    "    assert df.shape == X.shape and np.all(df == X)\n",
    "  \n",
    "  # reshape X such that flattened.shape = time, [x0, y0, z0, x1, y1, z1, ... xn-1, yn-1, zn-1]\n",
    "  flattened = X.swapaxes(0, 1).reshape( (df.shape[1], df.shape[0] * df.shape[2]), order='C' )\n",
    "\n",
    "  if run_tests: # switch to false to skip tests\n",
    "    for idx, i in enumerate(df):\n",
    "      for jdx, j in enumerate(df[idx]):\n",
    "        for kdx, k in enumerate(df[idx][jdx]):\n",
    "          assert flattened[jdx][ (idx*df.shape[2]) + kdx ] == df[idx][jdx][kdx]\n",
    "          \n",
    "  return flattened\n",
    "\n",
    "def unflatten(df, run_tests=True, start_time_index=0):\n",
    "  '''\n",
    "  df is a numpy array with the following two axes:\n",
    "    df.shape[0] = time index\n",
    "    df.shape[1] = [vertex_index*3] + dimension_vertex\n",
    "    \n",
    "  To unflatten this dataframe will mean to push the data into shape:\n",
    "    unflattened.shape[0] = the index of a vertex\n",
    "    unflattened.shape[1] = the index of a time stamp\n",
    "    unflattened.shape[2] = the index of a dimension (x, y, z)\n",
    "    \n",
    "  So df[2][4] == unflattened[1][2][0]\n",
    "  '''\n",
    "  if run_tests:\n",
    "    assert (len(df.shape) == 2) and (df.shape[1] == X.shape[0] * X.shape[2])\n",
    "  \n",
    "  unflattened = np.zeros(( X.shape[0], df.shape[0], X.shape[2] ))\n",
    "\n",
    "  for idx, i in enumerate(df):\n",
    "    for jdx, j in enumerate(df[idx]):\n",
    "      kdx = int(floor(jdx / 3))\n",
    "      ldx = int(jdx % 3)\n",
    "      unflattened[kdx][idx][ldx] = df[idx][jdx]\n",
    "\n",
    "  if run_tests: # set to false to skip tests\n",
    "    for idx, i in enumerate(unflattened):\n",
    "      for jdx, j in enumerate(unflattened[idx]):\n",
    "        for kdx, k in enumerate(unflattened[idx][jdx]):\n",
    "          assert( unflattened[idx][jdx][kdx] == X[idx][int(start_time_index)+jdx][kdx] )\n",
    "\n",
    "  return unflattened"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flat = flatten(X)\n",
    "unflat = unflatten(flat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mpl_toolkits.mplot3d.axes3d as p3\n",
    "from mpl_toolkits.mplot3d.art3d import juggle_axes\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import HTML\n",
    "from matplotlib import animation\n",
    "from copy import deepcopy\n",
    "import matplotlib\n",
    "\n",
    "# ask matplotlib to plot up to 2^128 frames in animations\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**128\n",
    "\n",
    "def update_points(time, points, df):\n",
    "  '''\n",
    "  Callback function called by plotting function below. Mutates the vertex\n",
    "  positions of each value in `points` so the animation moves\n",
    "  @param int time: the index of the time slice to visualize within `df`\n",
    "  @param mpl_toolkits.mplot3d.art3d.Path3DCollection points: the actual\n",
    "    geometry collection whose internal values this function mutates to move\n",
    "    the displayed points\n",
    "  @param numpy.ndarray df: a numpy array with the following three axes:\n",
    "    df.shape[0] = n_vertices\n",
    "    df.shape[1] = n_time_slices\n",
    "    df.shape[2] = n_dimensions\n",
    "  '''\n",
    "  points._offsets3d = juggle_axes(df[:,time,0], df[:,time,1], df[:,time,2], 'z')\n",
    "\n",
    "def get_plot(df, axis_min=0, axis_max=1, frames=200, speed=45, start_time_index=0, run_tests=True):\n",
    "  '''\n",
    "  General function that can plot numpy arrays in either of two shapes.\n",
    "  @param numpy.ndarray df: a numpy array with either of the following two shapes:\n",
    "    Possibility one:\n",
    "      df.shape[0] = n_vertices\n",
    "      df.shape[1] = n_time_slices\n",
    "      df.shape[2] = n_dimensions\n",
    "    Possibility two:\n",
    "      df.shape[0] = n_time_slices\n",
    "      df.shape[1] = [x0, y0, z0, x1, y1, z1, ... xn-1, yn-1, zn-1]\n",
    "    If the latter is received, we \"unflatten\" the df into the three dimensional variant\n",
    "  @param int axis_min: the minimum value of each axis scale\n",
    "  @param int axis_max: the maximum value of each axis scale\n",
    "  @param int frames: the number of time slices to animate.\n",
    "  @param int speed: the temporal duration of each frame. Increase to boost fps.\n",
    "  @param int start_time_index: the index position of the first frame in df within X. In other\n",
    "    words, if df starts at the nth time frame from X, start_time_index = n.\n",
    "  @param bool run_tests: boolean indicating whether we'll run the data validation\n",
    "    tests, should we need to unflatten the array. Should be set to False if we're passing\n",
    "    in predicted values, as they'll differ from X values.\n",
    "  '''\n",
    "  df = deepcopy(df)\n",
    "  if len(df.shape) == 2:\n",
    "    df = unflatten(df, start_time_index=start_time_index, run_tests=run_tests)\n",
    "  # center the data for visualization\n",
    "  df -= np.amin(df, axis=(0, 1))\n",
    "  df /= np.amax(df, axis=(0, 1))\n",
    "  # scale the z dimension\n",
    "  df[:,:,2] *= -1\n",
    "  df[:,:,2] += 1  \n",
    "  # plot the data\n",
    "  fig = plt.figure()\n",
    "  ax = p3.Axes3D(fig)\n",
    "  ax.set_xlim(axis_min, axis_max)\n",
    "  ax.set_ylim(axis_min, axis_max)\n",
    "  ax.set_zlim(axis_min, axis_max*1.5)\n",
    "  points = ax.scatter(df[:,0,0], df[:,0,1], df[:,0,2], depthshade=False) # x,y,z vals\n",
    "  return animation.FuncAnimation(fig,\n",
    "    update_points,\n",
    "    frames,\n",
    "    interval=speed,\n",
    "    fargs=(points, df),\n",
    "    blit=False  \n",
    "  ).to_html5_video()\n",
    "\n",
    "HTML(get_plot(unflat, frames=200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_x has shape: n_samples, look_back, n_vertices*3\n",
    "look_back = 128 # number of previous time slices to use to predict the time positions at time `i`\n",
    "train_x = []\n",
    "train_y = []\n",
    "\n",
    "# each i is a time slice; these time slices start at idx `look_back` (so we can look back `look_back` slices)\n",
    "for i in range(look_back, n_time-1, 1):\n",
    "  train_x.append( flat[i-look_back:i, :] )\n",
    "  train_y.append( flat[i:i+1] )\n",
    "  \n",
    "train_x = np.array(train_x)\n",
    "train_y = np.concatenate(train_y, axis=0)\n",
    "\n",
    "train_x.shape, train_y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visually confirm that the train_y is a valid dance sequence\n",
    "HTML(get_plot(train_y, frames=20, start_time_index=look_back))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visually confirm the train_x data is properly formatted - (only if look_back is 1)\n",
    "if look_back == 1:\n",
    "  HTML(get_plot(train_x.squeeze(), frames=20, start_time_index=look_back-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.mdn import MDN\n",
    "from keras.models import Sequential, Model\n",
    "from keras.layers import Dense, LSTM, Dropout, Activation, CuDNNLSTM\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras.losses import mean_squared_error\n",
    "from keras.optimizers import Adam\n",
    "from keras import backend as K\n",
    "import keras, os\n",
    "\n",
    "# config\n",
    "cells = [32, 32, 32, 32] # number of cells in each lstm layer\n",
    "output_dims = int(X.shape[0]) * int(X.shape[2]) # number of coordinate values to be predicted by each gaussian model\n",
    "input_shape = (look_back, output_dims) # shape of each input feature\n",
    "use_mdn = True # whether to use the MDN final layer or not\n",
    "n_mixes = 2 # number of gaussian models to build if use_mdn == True\n",
    "\n",
    "# optimizer params\n",
    "lr = 0.00001 # the learning rate of the model\n",
    "optimizer = Adam(lr=lr, clipvalue=0.5)\n",
    "\n",
    "# use tensorflow backend\n",
    "os.environ['KERAS_BACKEND'] = 'tensorflow'\n",
    "\n",
    "# determine the LSTM cells to use (hinges on whether GPU is available to keras)\n",
    "gpus = K.tensorflow_backend._get_available_gpus()\n",
    "LSTM_UNIT = CuDNNLSTM if len(gpus) > 0 else LSTM\n",
    "print('GPUs found:', gpus)\n",
    "\n",
    "# build the model\n",
    "model = Sequential()\n",
    "model.add(LSTM_UNIT(cells[0], return_sequences=True, input_shape=input_shape, ))\n",
    "model.add(LSTM_UNIT(cells[1], return_sequences=True, ))\n",
    "model.add(LSTM_UNIT(cells[2], ))\n",
    "model.add(Dense(cells[3]), )\n",
    "\n",
    "if use_mdn:\n",
    "  mdn = MDN(output_dims, n_mixes)\n",
    "  model.add(mdn)\n",
    "  model.compile(loss=mdn.get_loss_func(), optimizer=optimizer, metrics=['accuracy'])\n",
    "else:\n",
    "  model.add(Dense(output_dims, activation='tanh'))\n",
    "  model.compile(loss=mean_squared_error, optimizer=optimizer, metrics=['accuracy'])\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check untrained (baseline) accuracy\n",
    "samples = 1\n",
    "model.evaluate(train_x[:samples], train_y[:samples])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.callbacks import TerminateOnNaN\n",
    "from livelossplot import PlotLossesKeras\n",
    "from datetime import datetime\n",
    "import time, keras, os, json\n",
    "  \n",
    "class Logger(keras.callbacks.Callback):\n",
    "  '''Save the model and its weights every `self.save_frequency` epochs'''\n",
    "  def __init__(self):\n",
    "    self.epoch = 0 # stores number of completed epochs\n",
    "    self.save_frequency = 1 # configures how often we'll save the model and weights\n",
    "    self.date = datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d-%H:%M')\n",
    "    if not os.path.exists('snapshots'): os.makedirs('snapshots')\n",
    "    self.save_config()\n",
    "    \n",
    "  def save_config(self):\n",
    "    with open('snapshots/' + self.date + '-config.json', 'w') as out:\n",
    "      json.dump({\n",
    "        'look_back': look_back,\n",
    "        'cells': cells,\n",
    "        'use_mdn': use_mdn,\n",
    "        'n_mixes': n_mixes,\n",
    "        'lr': lr,\n",
    "      }, out)\n",
    "  \n",
    "  def on_batch_end(self, batch, logs={}, shape=train_x.shape):\n",
    "    if (batch+1 == shape[0]): # batch value is batch index, which is 0-based\n",
    "      self.epoch += 1\n",
    "      if (self.epoch > 0) and (self.epoch % self.save_frequency == 0):\n",
    "        path = 'snapshots/' + self.date + '-' + str(batch)\n",
    "        model.save(path + '.model')\n",
    "        model.save_weights(path + '.weights')\n",
    "\n",
    "#K.set_value(optimizer.lr, 0.00001)\n",
    "callbacks = [Logger(), TerminateOnNaN()]\n",
    "history = model.fit(train_x, train_y, epochs=50, batch_size=1, shuffle=False, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# grok the layer weights\n",
    "for layer_idx, layer in enumerate(model.layers): \n",
    "  for weight_block in layer.get_weights():\n",
    "    print(np.min(weight_block), np.max(weight_block), weight_block.shape, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save = False\n",
    "load = True\n",
    "model_path = 'snapshots/2018-10-26-17:53-496'\n",
    "if not os.path.exists('snapshots'): os.makedirs('snapshots')\n",
    "\n",
    "if save:\n",
    "  model.save(model_path + '.model')\n",
    "  model.save_weights(model_path + '.weights')\n",
    "  \n",
    "if load:\n",
    "  model.load_weights(model_path + '.weights')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize Results\n",
    "\n",
    "First let's analyze how well the model learned the input sequence\n",
    "\n",
    "### Assess Model Performance on Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize how well the model learned the input sequence\n",
    "n_frames = 500 # n frames of time slices to generate\n",
    "frames = []\n",
    "\n",
    "test_x = train_x[:n_frames] # data to pass into forward prop through the model\n",
    "y_pred = model.predict(test_x) # output with shape (n_frames, (output_dims+2) * n_mixes )\n",
    "\n",
    "# partition out the mus, sigs, and mixture weights\n",
    "for i in range(n_frames):\n",
    "  y = y_pred[i].squeeze()\n",
    "  mus = y[:n_mixes*output_dims]\n",
    "  sigs = y[n_mixes*output_dims:n_mixes*output_dims + n_mixes]\n",
    "  alphas = y[-n_mixes:]\n",
    "\n",
    "  # find the most likely distribution - then disregard that number and use the first Gaussian :)\n",
    "  alpha_idx = np.argmax(alphas)\n",
    "  alpha_idx = 0\n",
    "  \n",
    "  # pull out the mus that correspond to the selected alpha index\n",
    "  positions = mus[alpha_idx * output_dims:(alpha_idx+1) * output_dims]\n",
    "  frames.append(positions)\n",
    "  \n",
    "frames = np.array(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# skip the first look_back frames\n",
    "HTML(get_plot(frames, frames=n_frames, run_tests=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assess Model's Ability to Generate New Sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax(x):\n",
    "  \"\"\"Compute softmax values for each sets of scores in x.\"\"\"\n",
    "  r = np.exp(x - np.max(x))\n",
    "  return r / r.sum()\n",
    "\n",
    "n_frames = 500 # n frames of time slices to generate\n",
    "frames = []\n",
    "\n",
    "seed = np.random.randint(0, len(train_x)-1)\n",
    "x = np.expand_dims(train_x[seed], axis=0)\n",
    "print(' * seeding with', seed)\n",
    "\n",
    "for i in range(n_frames):\n",
    "  y = model.predict(x).squeeze()\n",
    "  mus = y[:n_mixes*output_dims]\n",
    "  sigs = y[n_mixes*output_dims:-n_mixes]\n",
    "  alphas = softmax(y[-n_mixes:])\n",
    "  \n",
    "  # select the alpha channel to use\n",
    "  alpha_idx = np.argmax(alphas)\n",
    "  \n",
    "  # grab the mus and sigs associated with the selected alpha_idx\n",
    "  frame_mus = mus.ravel()[alpha_idx*output_dims : (alpha_idx+1)*output_dims]\n",
    "  frame_sig = sigs[alpha_idx] / 100\n",
    "  \n",
    "  # now sample from each Gaussian\n",
    "  positions = [np.random.normal(loc=m, scale=frame_sig) for m in frame_mus]\n",
    "  positions = frame_mus\n",
    "  \n",
    "  # add these positions to the results\n",
    "  frames.append(positions)\n",
    "  \n",
    "  # pull out a new training example - stack the new result on\n",
    "  # all values after the first from the bottom-most value in the x's\n",
    "  start = x[:,1:,:]\n",
    "  end = np.expand_dims( np.expand_dims(positions, axis=0), axis=0 )\n",
    "  x = np.concatenate((start, end), axis=1)\n",
    "  \n",
    "frames = np.array(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML(get_plot(frames, frames=n_frames, run_tests=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the new frames as a JSON array for WebGL visualizations\n",
    "\n",
    "import json\n",
    "\n",
    "with open('dance.json', 'w') as out:\n",
    "  json.dump(unflatten(frames, run_tests=False).tolist(), out)"
   ]
  }
 ],
 "metadata": {
  "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}