Code2Pix.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DSL2Pix - Deep Learning Compiler for User Interfaces\n",
    "---\n",
    "##### Noah Gundotra |  SUSA @ UC Berkeley\n",
    "This work was made possible by the coordination of UC Berkeley's Statistics Undergraduate Student Association [SUSA](https://susa.berkeley.edu/) and [Uizard](https://uizard.io/). This notebook is research sponsored by Uizard. \n",
    "\n",
    "This notebook encapsulates the model design process to create a differentiable compiler that can be used to effectively compile simple user interfaces for Android, iOS, and Web backends. Specifically, we are \"compiling\" from a domain specific language (DSL) introduced by Tony Beltramelli in his original [pix2code paper](https://arxiv.org/abs/1705.07962). In fact, the dataset and vocabulary assets come from the [public pix2code repository](https://github.com/tonybeltramelli/pix2code).\n",
    "\n",
    "The successful model design ended up being using 1D convolutional layers for analyzing the DSL tokens, paired with a convolutional decoder taken from autoencoders trained on the pix2code datasets. See the Hydra Autoencoder.ipynb to train these.\n",
    "\n",
    "In the near future, we hope to use this model to improve the accuracy of pix2code and related models by using dsl2pix in a GAN with pix2code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "__authors__ = 'Noah Gundotra, Samyak Parajuli, Luke Dai, Ajay Raj, Aismit Das, Dennis Yang, Japjot Singh'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-04-18T20:37:43.261847Z",
     "start_time": "2018-04-18T20:37:39.755871Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt \n",
    "% matplotlib inline\n",
    "import os\n",
    "from os.path import join\n",
    "\n",
    "import keras\n",
    "from keras.models import Sequential\n",
    "from keras.layers import (LSTM, Dense, Conv2D, Input, Reshape, concatenate, MaxPooling2D, Dropout, Flatten, \n",
    "RepeatVector, UpSampling2D, Conv1D, Permute, BatchNormalization, Activation, \n",
    "                          UpSampling2D, MaxPooling1D, GlobalAveragePooling1D, Embedding)\n",
    "from keras import Model\n",
    "from keras.engine.topology import Layer, InputSpec\n",
    "from keras import backend as K\n",
    "from keras.optimizers import RMSprop\n",
    "from keras.callbacks import EarlyStopping\n",
    "from keras.utils import to_categorical\n",
    "from keras.losses import mean_squared_error\n",
    "from keras.preprocessing import sequence\n",
    "\n",
    "from Vocabulary import Vocabulary\n",
    "from keras.backend import clear_session\n",
    "\n",
    "clear_session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose your backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-04-18T20:37:43.267061Z",
     "start_time": "2018-04-18T20:37:43.263712Z"
    }
   },
   "outputs": [],
   "source": [
    "datasets = ['android', 'ios', 'web']\n",
    "BACKEND = datasets[2]\n",
    "traindir = \"{}/training_features\".format(BACKEND)\n",
    "vocab_path = '../pix2code/bin/{}/'.format(BACKEND)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data (tokens, images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_string(data_dir, file):\n",
    "    \"\"\"Turns the gui descriptor files into a list of their 'words'\"\"\"\n",
    "    f = open(join(data_dir, file))\n",
    "    original_string = f.read()\n",
    "    list_words = original_string.split()\n",
    "    new_list = []\n",
    "    for word in list_words:\n",
    "        # Remove special characters\n",
    "        special_characters = \" ,\\n\"\n",
    "        for c in special_characters:\n",
    "            word = word.replace(c, '')\n",
    "        new_list.append(word)\n",
    "    return new_list, original_string\n",
    "\n",
    "def load_data(data_dir):\n",
    "    \"\"\"Loads in the sentence and image data together.\n",
    "    Returns (sentences, images, original strings (in descriptor files))\"\"\"\n",
    "    images = []\n",
    "    sentences = []\n",
    "    original_strings = []\n",
    "    files = os.listdir(data_dir)\n",
    "    for listing in files:\n",
    "        # Load the sentence files, then the images\n",
    "        # There are less likely to be user-made files in this dir that end with .gui\n",
    "        if listing.endswith('.gui'):\n",
    "            base = listing[:-4]\n",
    "            try:\n",
    "                img_f = join(data_dir, base + '.npz')\n",
    "                img = np.load(img_f)['features']\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                print(\"Error trying to match img for {}\".format(listing))\n",
    "                continue\n",
    "            sent, original_string = process_string(data_dir, listing)\n",
    "            sentences.append(sent)\n",
    "            original_strings.append(original_string)\n",
    "            images.append(img)\n",
    "    return sentences, np.array(images), original_strings\n",
    "\n",
    "sentences, GUIS, original_strings = load_data(traindir)\n",
    "assert len(sentences) == len(GUIS) and len(sentences) == 1500\n",
    "print(\"Loaded succesfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# One Hot\n",
    "\n",
    "Here we are drawing on files and a class from pix2code training. This Vocabulary class is from `pix2code/model/classes/Vocabulary.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voc = Vocabulary()\n",
    "voc.retrieve(vocab_path)\n",
    "VOCAB_SIZE = len(voc.binary_vocabulary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-04-18T20:38:24.701841Z",
     "start_time": "2018-04-18T20:38:23.574258Z"
    }
   },
   "outputs": [],
   "source": [
    "#Find max seq length\n",
    "max_length = 0\n",
    "for j in range(len(sentences)):\n",
    "    if len(sentences) > max_length:\n",
    "        max_length = len(sentences)\n",
    "\n",
    "def to_hot(sentences, max_length, voc):\n",
    "    hot = np.zeros((len(sentences), max_length, len(voc.binary_vocabulary)))\n",
    "    for i, sent in enumerate(sentences):\n",
    "        for j, word in enumerate(sent):\n",
    "            hot[i, j] = voc.binary_vocabulary[word]\n",
    "    return hot\n",
    "\n",
    "hot = to_hot(sentences, max_length, voc)\n",
    "assert len(hot) == 1500\n",
    "print(\"Converted to one-hot succesfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Design\n",
    "\n",
    "Functions that help the model design process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_frozen(model_path, freeze=False):\n",
    "    \"\"\"Loads in the model from model_path. If freeze=True, then the model will be non-trainable.\"\"\"\n",
    "    model = keras.models.load_model(model_path)\n",
    "    for layer in model.layers:\n",
    "        layer.trainable = False if freeze else True\n",
    "    return model\n",
    "\n",
    "def make_callback(model, dataset):\n",
    "    \"\"\"Creates a lambda callback that plots the models prediction on the first 20 pictures of the dataset.\"\"\"\n",
    "    def callback(epoch, logs):\n",
    "        gif_range(model, dataset, start=0, stop=20)\n",
    "    return keras.callbacks.LambdaCallback(on_epoch_end=callback)\n",
    "\n",
    "def shuffle_weights(model, weights=None):\n",
    "    \"\"\"\n",
    "    @author: jklient\n",
    "    @source: https://github.com/keras-team/keras/issues/341\n",
    "\n",
    "    Randomly permute the weights in `model`, or the given `weights`.\n",
    "\n",
    "    This is a fast approximation of re-initializing the weights of a model.\n",
    "\n",
    "    Assumes weights are distributed independently of the dimensions of the weight tensors\n",
    "      (i.e., the weights have the same distribution along each dimension).\n",
    "\n",
    "    :param Model model: Modify the weights of the given model.\n",
    "    :param list(ndarray) weights: The model's weights will be replaced by a random permutation of these weights.\n",
    "      If `None`, permute the model's current weights.\n",
    "    \"\"\"\n",
    "    if weights is None:\n",
    "        weights = model.get_weights()\n",
    "    weights = [np.random.permutation(w.flat).reshape(w.shape) for w in weights]\n",
    "    # Faster, but less random: only permutes along the first dimension\n",
    "    # weights = [np.random.permutation(w) for w in weights]\n",
    "    model.set_weights(weights)\n",
    "    \n",
    "def gif_range(model, data, start=0, stop=1, save=True):\n",
    "    \"\"\"\n",
    "    Adapted from - \n",
    "    @author: Eli Bendersky\n",
    "    @source: https://eli.thegreenplace.net/2016/drawing-animated-gifs-with-matplotlib/\n",
    "    Makes a GIF that reveals the model's predictions for the first (stop-start) images in data.\n",
    "    Saves the GIF to `dsl_predictions.gif`\n",
    "    \"\"\"\n",
    "    import sys\n",
    "    from matplotlib.animation import FuncAnimation\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    #   fig.set_tight_layout(True)\n",
    "\n",
    "    # Query the figure's on-screen size and DPI. Note that when saving the figure to\n",
    "    # a file, we need to provide a DPI for that separately.\n",
    "    #   print('fig size: {0} DPI, size in inches {1}'.format(\n",
    "    #       fig.get_dpi(), fig.get_size_inches()))\n",
    "    \n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "        \n",
    "    def update(i):\n",
    "        label = 'wordembedding {0}'.format(i)\n",
    "        # print(label)\n",
    "        ax.imshow(model.predict(np.expand_dims(data[i], axis=0))[0])\n",
    "        ax.set_xlabel(label)\n",
    "        return ax\n",
    "    \n",
    "    fig.tight_layout()\n",
    "    anim = FuncAnimation(fig, update, frames=np.arange(start, stop), interval=70)\n",
    "    if save:\n",
    "        anim.save('dsl_predictions.gif', dpi=80, writer='imagemagick')\n",
    "    else:\n",
    "        # plt.show() will just loop the animation forever.\n",
    "        plt.show()\n",
    "    plt.close(fig)\n",
    "    \n",
    "def unroll_hydra(model, last):\n",
    "    \"\"\"Unrolls the decoder head of the hydra model given. Each hydra model has only 1 head, so this\n",
    "    should be relatively simple.\"\"\"\n",
    "    return model.layers[-1](last)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Code2Pix\n",
    "\n",
    "The following code cells create the code2pix model and train it.\n",
    "\n",
    "\n",
    "## Pro-Tip\n",
    "\n",
    "Instead of just looking at numbers while this trains - look at pretty pictures! The `val_plot` callback to this model will update `dsl_predictions.gif` with the model's predictions on the first 20 images in the validation data. It's pretty cool to watch. **Highly recommend.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_conv(autoencoder_path):\n",
    "    \"\"\"1D conv model :/\"\"\"\n",
    "    x_in = Input(shape=(max_length, VOCAB_SIZE))\n",
    "    last = Conv1D(64, 7, padding='same', activation='relu')(x_in)\n",
    "    last = MaxPooling1D()(last)\n",
    "    last = Conv1D(64, 5, padding='same', activation='relu')(last)\n",
    "    last = Conv1D(128, 4, padding='same', activation='relu')(last)\n",
    "    last = MaxPooling1D()(last)\n",
    "    last = Flatten()(last)\n",
    "    last = Dense(1024, activation='relu')(last)\n",
    "    last = Reshape((8, 8, 16))(last)\n",
    "    last = UpSampling2D()(last)\n",
    "    autoencoder = load_frozen(autoencoder_path, freeze=False)\n",
    "    last = unroll_hydra(autoencoder, last)\n",
    "    model = Model(x_in, last)\n",
    "    model.compile(RMSprop(lr=0.01), loss='binary_crossentropy')\n",
    "    return model\n",
    "\n",
    "clear_session()\n",
    "model = make_conv('{}-hydra-model'.format(BACKEND))\n",
    "\n",
    "# Uncomment to see the number of model parameters\n",
    "# model.summary()\n",
    "\n",
    "val_plot = make_callback(model, hot[1400:])\n",
    "\n",
    "# Early Stopping to prevent severe overfitting\n",
    "early_stop = EarlyStopping(monitor='val_loss', patience=7)\n",
    "\n",
    "model.fit(hot[:1400], GUIS[:1400], callbacks=[val_plot, early_stop],\n",
    "          validation_data=(hot[1400:], GUIS[1400:]), epochs=65, batch_size=32, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save('dsl2pix-{}'.format(BACKEND))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code cell chooses a random picture from the validation set and plots the input and output of the code2pix model and compares it to the ground-truth output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "def compare(sentences, hot, pics, idx):\n",
    "    \"\"\"Compares the results of the model.\n",
    "\n",
    "    Displays the input to the model, output of the model on the input, and ground truth.\n",
    "    Input to the model is chosen as the hot[idx].\n",
    "    \"\"\"\n",
    "    fig, axes = plt.subplots(1, 2)\n",
    "    for ax in axes:\n",
    "        ax.set_yticks([])\n",
    "        ax.set_xticks([])\n",
    "        \n",
    "    axes[0].imshow(model.predict(np.expand_dims(hot[idx], 0))[0])\n",
    "    axes[0].set_xlabel(\"Predicted GUI\")\n",
    "    axes[1].imshow(pics[idx])\n",
    "    axes[1].set_xlabel(\"Actual GUI\")\n",
    "    \n",
    "    print('gui descriptor:\\t', ' '.join(sentences[idx]))\n",
    "    count = Counter(sentences[idx])\n",
    "    print(count)\n",
    "    fig.savefig(\"{}-code2pix-predict.png\".format(BACKEND))\n",
    "    return count\n",
    "\n",
    "count = compare(sentences, hot, GUIS, np.random.choice(list(range(1400, 1500))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code cells are left for additional experimentation and educational purposes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Word2Vec Model - For fine tuned embeddings\n",
    "---\n",
    "\n",
    "Embedding the DSL sentence representation into condensed, semantically-infused vector representations. **Currently unused.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-04-18T20:38:27.084892Z",
     "start_time": "2018-04-18T20:38:26.107041Z"
    }
   },
   "outputs": [],
   "source": [
    "#All word embeddings (size of training set x sentence length x 100(embedding dimension))\n",
    "def create_wordembeddings(sentences, EMBED_SIZE):\n",
    "    \"\"\"Creates a Word2Vec model to compress the sentence data.\n",
    "    Returns a np.array of (#sentences, sentence length, EMBED_SIZE)\"\"\"\n",
    "    from gensim.models import Word2Vec\n",
    "    \n",
    "    model = Word2Vec(sentences, min_count=1,size=EMBED_SIZE)\n",
    "    words = list(model.wv.vocab)\n",
    "    \n",
    "    wordembeddings = []\n",
    "    for sentence in sentences:\n",
    "        result_array = np.empty((0, EMBED_SIZE))\n",
    "        for word in sentence:\n",
    "            result = model[word].reshape(1, EMBED_SIZE)\n",
    "            result_array = np.append(result_array, result, axis=0)\n",
    "        result_array = np.pad(result_array, ((max_length - result_array.shape[0],0), (0,0)), 'constant', constant_values = 0)\n",
    "        wordembeddings.append(result_array)\n",
    "    wordembeddings = np.array(wordembeddings)\n",
    "    return wordembeddings\n",
    "\n",
    "wordembeddings = create_wordembeddings(sentences, VOCAB_SIZE)\n",
    "assert wordembeddings.shape[0] == 1500"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extra Model Parts\n",
    "\n",
    "**Currently unused.** Used in previous iterations to repeat 3D embedding data. For example repeating a (16, 16, 16) embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-04-18T20:38:29.800684Z",
     "start_time": "2018-04-18T20:38:29.769226Z"
    }
   },
   "outputs": [],
   "source": [
    "#Custom Layer to get from 3D -> 4D\n",
    "class RepeatVector4D(Layer):\n",
    "    def __init__(self, n, **kwargs):\n",
    "        self.n = n\n",
    "        self.input_spec = [InputSpec(ndim=3)]\n",
    "        super(RepeatVector4D, self).__init__(**kwargs)\n",
    "        \n",
    "    def get_output_shape_for(self, input_shape):\n",
    "        return (input_shape[0], self.n, input_shape[1], input_shape[2])\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return (input_shape[0], self.n, input_shape[1], input_shape[2])\n",
    "\n",
    "    def call(self, x, mask=None):\n",
    "        x = K.expand_dims(x, 1)\n",
    "        pattern = K.stack([1, self.n, 1, 1])\n",
    "        return K.tile(x, pattern)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM-Based Model\n",
    "\n",
    "This model is based off the pix2code model design. This model uses lstms to encode the token data, and loads in pretrained lstms from the pix2code models. Super slow to train ~3min/epoch. ***Deprecated.*** If you would like to run these, feel free to, but the rest of the code is not explicitly designed to support these models anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_model(encoder_path, lstm_decoder_path, autoencoder_path):\n",
    "    # Freeze the pretrained encoder & decoder (respectively)\n",
    "    pix2code_lstm = load_frozen(encoder_path, False)\n",
    "    lstm_decoder = load_frozen(lstm_decoder_path, False)\n",
    "    autoencoder = load_frozen(autoencoder_path, False)\n",
    "\n",
    "    #Input DSL -> Intermediate Encoding through convolution\n",
    "    x_in = Input(shape = (max_length, 19), name = 'x_in')\n",
    "    drop_rate = 0.01\n",
    "#     last = Dropout(drop_rate)(last)\n",
    "#     #y_lstm = LSTM(128, return_sequences = True, input_shape=(max_length, 100))(y_in)\n",
    "    \n",
    "#     b1 = Dense(48*48, activation='relu')(x_flatten)\n",
    "#     b1 = Reshape((48, 48, 1))(b1)\n",
    "#     b1 = Conv2D(filters=64, kernel_size=(3,3), strides=(3,3), padding='valid', \n",
    "#                 activation='relu', name='b1conv_1')(b1)\n",
    "#     b1 = Conv2D(16, (3,3), strides=(1,1), padding='same', activation='relu')(b1)\n",
    "#     last = Reshape((48, 19))(last) # Input size of pix2code's first encoder\n",
    "    last = pix2code_lstm(x_in)\n",
    "    last = LSTM(512, return_sequences = True)(last)\n",
    "    last = Dropout(drop_rate)(last)\n",
    "    last = lstm_decoder.layers[-1](last)\n",
    "    reshape = Reshape((8,8,8))(last)\n",
    "    last = UpSampling2D((2,2), name='upsampler-trainable')(reshape)\n",
    "    last = Conv2D(16, kernel_size=(3,3), padding='same', activation='relu')(last)\n",
    "    last = Reshape((16,16,16))(last)\n",
    "\n",
    "    # Load in the v0 autoencoder\n",
    "#     last = keras.layers.Add()([b1, last])\n",
    "#     last = Conv2D(16, kernel_size=(2,2), padding='same',activation='relu')(last)\n",
    "    last = unroll_hydra(autoencoder, last)\n",
    "    \n",
    "    model = Model(x_in, last)\n",
    "    opt = RMSprop(lr=0.001)\n",
    "    model.compile(optimizer=opt, loss='binary_crossentropy')\n",
    "    return model\n",
    "\n",
    "clear_session()\n",
    "model = make_model('ios-p2c-encoder-LSTM', 'ios-p2c-decoder-LSTM', 'ios-hydra-model')\n",
    "# model.summary()\n",
    "callback = make_callback(model, hot)\n",
    "# shuffle_weights(model)\n",
    "model.fit(hot, GUIS, callbacks=[callback], epochs=200, batch_size=32, validation_split=0.2, shuffle=False)\n",
    "\n",
    "\n",
    "def make_model_old(encoder_path, lstm_decoder_path, autoencoder_path='model/autoencoders-v0/ios.h5'):\n",
    "    \"\"\"A deprecated iteration of the dsl2pix model design.\"\"\"\n",
    "    # Freeze the pretrained encoder & decoder (respectively)\n",
    "    pix2code_lstm = load_frozen(encoder_path, False)\n",
    "    autoencoder = load_frozen(autoencoder_path, False)\n",
    "\n",
    "    x_in = Input(shape = (max_length, VOCAB_SIZE), name = 'x_in')\n",
    "    x_flatten = Flatten()(x_in)\n",
    "    last = Dense(912, activation='relu')(x_flatten)\n",
    "    drop_rate = 0.1\n",
    "    last = Dropout(drop_rate)(last)\n",
    "    #y_lstm = LSTM(128, return_sequences = True, input_shape=(max_length, 100))(y_in)\n",
    "    \n",
    "    b1 = Dense(48*48, activation='relu')(x_flatten)\n",
    "    b1 = Reshape((48, 48, 1))(b1)\n",
    "    b1 = Conv2D(filters=64, kernel_size=(3,3), strides=(3,3), padding='valid', \n",
    "                activation='relu', name='b1conv_1')(b1)\n",
    "    b1 = Conv2D(16, (3,3), strides=(1,1), padding='same', activation='relu')(b1)\n",
    "    last = Reshape((48, 19))(last) # Input size of pix2code's first encoder\n",
    "    last = LSTM(19, return_sequences=True)(last)\n",
    "    last = LSTM(19, return_sequences=True)(last)\n",
    "    last = pix2code_lstm(last)\n",
    "    last = LSTM(128, return_sequences = True)(last)\n",
    "    last = LSTM(256, return_sequences = True)(last)\n",
    "    last = Dropout(drop_rate)(last)\n",
    "    last = LSTM(512, return_sequences = False)(last)\n",
    "    reshape = Reshape((8,8,8))(last)\n",
    "    last = UpSampling2D((2,2), name='upsampler-trainable')(reshape)\n",
    "    last = Conv2D(32, kernel_size=(4,4), padding='same', activation='relu')(last)\n",
    "#     last = Dropout(drop_rate)(last)\n",
    "    last = Conv2D(16, kernel_size=(3,3), padding='same', activation='relu')(last)\n",
    "    last = Reshape((16,16,16))(last)\n",
    "    \n",
    "    # Load in the v0 autoencoder\n",
    "    last = keras.layers.Add()([b1, last])\n",
    "    last = Conv2D(16, kernel_size=(2,2), padding='same',activation='relu')(last)\n",
    "    last = unroll_v1(autoencoder, last)\n",
    "    \n",
    "    model = Model(x_in, last)\n",
    "    model.compile(optimizer='rmsprop', loss='binary_crossentropy')\n",
    "    return model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch Kernel",
   "language": "python",
   "name": "torchkern"
  },
  "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}