CapsuleNetMNIST-tf2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "language": "python",
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "version": 3,
        "name": "ipython"
      },
      "file_extension": ".py",
      "pygments_lexer": "ipython3",
      "mimetype": "text/x-python",
      "nbconvert_exporter": "python",
      "version": "3.6.3",
      "name": "python"
    },
    "colab": {
      "name": "CapsuleNet on MNIST.ipynb",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/bklooste/tensorflowcollab/blob/master/CapsuleNetMNIST-tf2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "_cell_guid": "7f4c4917-22fb-4b10-879e-51c600e6c3af",
        "_uuid": "9b086e5ec535ad75eada3ca72bf5e6534251074f",
        "id": "LQnu9movo8xz",
        "colab_type": "text"
      },
      "source": [
        "# Overview\n",
        "\n",
        "The new model CapsuleNet proposed by Sara Sabour (and Geoffry Hinton) claims to deliver state of the art results on [MNIST](https://arxiv.org/abs/1710.09829). The kernel aims to create and train the model using the Kaggle Dataset and then make a submission to see where it actually ends up. Given the constraint of using a Kaggle Kernel means it can't be trained as long as we would like or with GPU's but IMHO if a model can't be reasonably well trained in an hour on a 28x28 dataset, that model probably won't be too useful in the immediate future.\n",
        "\n",
        "## Implementation Details\n",
        "\n",
        "* Keras implementation of CapsNet in Hinton's paper Dynamic Routing Between Capsules.\n",
        "* Code adapted from https://github.com/XifengGuo/CapsNet-Keras/blob/master/capsulenet.py\n",
        "*  Author: Xifeng Guo, E-mail: `guoxifeng1990@163.com`, Github: `https://github.com/XifengGuo/CapsNet-Keras`\n",
        "*     The current version maybe only works for TensorFlow backend. Actually it will be straightforward to re-write to TF code.\n",
        "*     Adopting to other backends should be easy, but I have not tested this. \n",
        "\n",
        "Result:\n",
        "    Validation accuracy > 99.5% after 20 epochs. Still under-fitting.\n",
        "    About 110 seconds per epoch on a single GTX1070 GPU card\n",
        "    \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "e30e2e10-a909-485d-be9d-dc6f592911a7",
        "_uuid": "c7e569699c6d067cd9fdf9c77299775e399b2ef3",
        "id": "qWnuqFe2o8x2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import os\n",
        "import pandas as pd\n",
        "#from keras.preprocessing.image import ImageDataGenerator\n",
        "#from keras import callbacks\n",
        "#from keras.utils.vis_utils import plot_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "_cell_guid": "9776fd57-44e0-4211-a7a5-c7e647a10704",
        "_uuid": "f4b5499a472b312d5c5f0274ad429567aced6841",
        "id": "xYHeLw1do8x5",
        "colab_type": "text"
      },
      "source": [
        "# Capsule Layers \n",
        "Here is the implementation of the necessary layers for the CapsuleNet. These are not optimized yet and can be made significantly more performant. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "167d39ca-ee32-4eec-a83b-86194252b14f",
        "_uuid": "90c180a9a8c20e3fb8a93c3eb42588927cfcd6b6",
        "id": "dQti96GTo8x5",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 71
        },
        "outputId": "4ecd1c7d-a003-4b57-f918-d8ec263da1dd"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.keras.backend as K\n",
        "from tensorflow.keras import initializers, layers\n",
        "\n",
        "print(tf.__version__)\n",
        "\n",
        "\n",
        "class Length(layers.Layer):\n",
        "    \"\"\"\n",
        "    Compute the length of vectors. This is used to compute a Tensor that has the same shape with y_true in margin_loss.\n",
        "    Using this layer as model's output can directly predict labels by using `y_pred = np.argmax(model.predict(x), 1)`\n",
        "    inputs: shape=[None, num_vectors, dim_vector]\n",
        "    output: shape=[None, num_vectors]\n",
        "    \"\"\"\n",
        "    def call(self, inputs, **kwargs):\n",
        "        return tf.sqrt(tf.reduce_sum(tf.square(inputs), -1) + K.epsilon())\n",
        "\n",
        "    def compute_output_shape(self, input_shape):\n",
        "        return input_shape[:-1]\n",
        "\n",
        "    def get_config(self):\n",
        "        config = super(Length, self).get_config()\n",
        "        return config\n",
        "\n",
        "\n",
        "class Mask(layers.Layer):\n",
        "    \"\"\"\n",
        "    Mask a Tensor with shape=[None, num_capsule, dim_vector] either by the capsule with max length or by an additional \n",
        "    input mask. Except the max-length capsule (or specified capsule), all vectors are masked to zeros. Then flatten the\n",
        "    masked Tensor.\n",
        "    For example:\n",
        "        ```\n",
        "        x = keras.layers.Input(shape=[8, 3, 2])  # batch_size=8, each sample contains 3 capsules with dim_vector=2\n",
        "        y = keras.layers.Input(shape=[8, 3])  # True labels. 8 samples, 3 classes, one-hot coding.\n",
        "        out = Mask()(x)  # out.shape=[8, 6]\n",
        "        # or\n",
        "        out2 = Mask()([x, y])  # out2.shape=[8,6]. Masked with true labels y. Of course y can also be manipulated.\n",
        "        ```\n",
        "    \"\"\"\n",
        "    def call(self, inputs, **kwargs):\n",
        "        if type(inputs) is list:  # true label is provided with shape = [None, n_classes], i.e. one-hot code.\n",
        "            assert len(inputs) == 2\n",
        "            inputs, mask = inputs\n",
        "        else:  # if no true label, mask by the max length of capsules. Mainly used for prediction\n",
        "            # compute lengths of capsules\n",
        "            x = tf.sqrt(tf.reduce_sum(tf.square(inputs), -1))\n",
        "            # generate the mask which is a one-hot code.\n",
        "            # mask.shape=[None, n_classes]=[None, num_capsule]\n",
        "            mask = tf.one_hot(indices=tf.argmax(x, 1), depth=x.shape[1])\n",
        "\n",
        "        # inputs.shape=[None, num_capsule, dim_capsule]\n",
        "        # mask.shape=[None, num_capsule]\n",
        "        # masked.shape=[None, num_capsule * dim_capsule]\n",
        "        masked = K.batch_flatten(inputs * tf.expand_dims(mask, -1))\n",
        "        return masked\n",
        "\n",
        "    def compute_output_shape(self, input_shape):\n",
        "        if type(input_shape[0]) is tuple:  # true label provided\n",
        "            return tuple([None, input_shape[0][1] * input_shape[0][2]])\n",
        "        else:  # no true label provided\n",
        "            return tuple([None, input_shape[1] * input_shape[2]])\n",
        "\n",
        "    def get_config(self):\n",
        "        config = super(Mask, self).get_config()\n",
        "        return config\n",
        "\n",
        "\n",
        "def squash(vectors, axis=-1):\n",
        "    \"\"\"\n",
        "    The non-linear activation used in Capsule. It drives the length of a large vector to near 1 and small vector to 0\n",
        "    :param vectors: some vectors to be squashed, N-dim tensor\n",
        "    :param axis: the axis to squash\n",
        "    :return: a Tensor with same shape as input vectors\n",
        "    \"\"\"\n",
        "    s_squared_norm = tf.reduce_sum(tf.square(vectors), axis, keepdims=True)\n",
        "    scale = s_squared_norm / (1 + s_squared_norm) / tf.sqrt(s_squared_norm + K.epsilon())\n",
        "    return scale * vectors\n",
        "\n",
        "\n",
        "class CapsuleLayer(layers.Layer):\n",
        "    \"\"\"\n",
        "    The capsule layer. It is similar to Dense layer. Dense layer has `in_num` inputs, each is a scalar, the output of the\n",
        "    neuron from the former layer, and it has `out_num` output neurons. CapsuleLayer just expand the output of the neuron\n",
        "    from scalar to vector. So its input shape = [None, input_num_capsule, input_dim_capsule] and output shape = \\\n",
        "    [None, num_capsule, dim_capsule]. For Dense Layer, input_dim_capsule = dim_capsule = 1.\n",
        "    :param num_capsule: number of capsules in this layer\n",
        "    :param dim_capsule: dimension of the output vectors of the capsules in this layer\n",
        "    :param routings: number of iterations for the routing algorithm\n",
        "    \"\"\"\n",
        "    def __init__(self, num_capsule, dim_capsule, routings=3,\n",
        "                 kernel_initializer='glorot_uniform',\n",
        "                 **kwargs):\n",
        "        super(CapsuleLayer, self).__init__(**kwargs)\n",
        "        self.num_capsule = num_capsule\n",
        "        self.dim_capsule = dim_capsule\n",
        "        self.routings = routings\n",
        "        self.kernel_initializer = initializers.get(kernel_initializer)\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        assert len(input_shape) >= 3, \"The input Tensor should have shape=[None, input_num_capsule, input_dim_capsule]\"\n",
        "        self.input_num_capsule = input_shape[1]\n",
        "        self.input_dim_capsule = input_shape[2]\n",
        "\n",
        "        # Transform matrix, from each input capsule to each output capsule, there's a unique weight as in Dense layer.\n",
        "        self.W = self.add_weight(shape=[self.num_capsule, self.input_num_capsule,\n",
        "                                        self.dim_capsule, self.input_dim_capsule],\n",
        "                                 initializer=self.kernel_initializer,\n",
        "                                 name='W')\n",
        "\n",
        "        self.built = True\n",
        "\n",
        "    def call(self, inputs, training=None):\n",
        "        # inputs.shape=[None, input_num_capsule, input_dim_capsule]\n",
        "        # inputs_expand.shape=[None, 1, input_num_capsule, input_dim_capsule, 1]\n",
        "        inputs_expand = tf.expand_dims(tf.expand_dims(inputs, 1), -1)\n",
        "\n",
        "        # Replicate num_capsule dimension to prepare being multiplied by W\n",
        "        # inputs_tiled.shape=[None, num_capsule, input_num_capsule, input_dim_capsule, 1]\n",
        "        inputs_tiled = tf.tile(inputs_expand, [1, self.num_capsule, 1, 1, 1])\n",
        "\n",
        "        # Compute `inputs * W` by scanning inputs_tiled on dimension 0.\n",
        "        # W.shape=[num_capsule, input_num_capsule, dim_capsule, input_dim_capsule]\n",
        "        # x.shape=[num_capsule, input_num_capsule, input_dim_capsule, 1]\n",
        "        # Regard the first two dimensions as `batch` dimension, then\n",
        "        # matmul(W, x): [..., dim_capsule, input_dim_capsule] x [..., input_dim_capsule, 1] -> [..., dim_capsule, 1].\n",
        "        # inputs_hat.shape = [None, num_capsule, input_num_capsule, dim_capsule]\n",
        "        inputs_hat = tf.squeeze(tf.map_fn(lambda x: tf.matmul(self.W, x), elems=inputs_tiled))\n",
        "\n",
        "        # Begin: Routing algorithm ---------------------------------------------------------------------#\n",
        "        # The prior for coupling coefficient, initialized as zeros.\n",
        "        # b.shape = [None, self.num_capsule, 1, self.input_num_capsule].\n",
        "        b = tf.zeros(shape=[inputs.shape[0], self.num_capsule, 1, self.input_num_capsule])\n",
        "\n",
        "        assert self.routings > 0, 'The routings should be > 0.'\n",
        "        for i in range(self.routings):\n",
        "            # c.shape=[batch_size, num_capsule, 1, input_num_capsule]\n",
        "            c = tf.nn.softmax(b, axis=1)\n",
        "\n",
        "            # c.shape = [batch_size, num_capsule, 1, input_num_capsule]\n",
        "            # inputs_hat.shape=[None, num_capsule, input_num_capsule, dim_capsule]\n",
        "            # The first two dimensions as `batch` dimension,\n",
        "            # then matmal: [..., 1, input_num_capsule] x [..., input_num_capsule, dim_capsule] -> [..., 1, dim_capsule].\n",
        "            # outputs.shape=[None, num_capsule, 1, dim_capsule]\n",
        "            outputs = squash(tf.matmul(c, inputs_hat))  # [None, 10, 1, 16]\n",
        "\n",
        "            if i < self.routings - 1:\n",
        "                # outputs.shape =  [None, num_capsule, 1, dim_capsule]\n",
        "                # inputs_hat.shape=[None, num_capsule, input_num_capsule, dim_capsule]\n",
        "                # The first two dimensions as `batch` dimension, then\n",
        "                # matmal:[..., 1, dim_capsule] x [..., input_num_capsule, dim_capsule]^T -> [..., 1, input_num_capsule].\n",
        "                # b.shape=[batch_size, num_capsule, 1, input_num_capsule]\n",
        "                b += tf.matmul(outputs, inputs_hat, transpose_b=True)\n",
        "        # End: Routing algorithm -----------------------------------------------------------------------#\n",
        "\n",
        "        return tf.squeeze(outputs)\n",
        "\n",
        "    def compute_output_shape(self, input_shape):\n",
        "        return tuple([None, self.num_capsule, self.dim_capsule])\n",
        "\n",
        "    def get_config(self):\n",
        "        config = {\n",
        "            'num_capsule': self.num_capsule,\n",
        "            'dim_capsule': self.dim_capsule,\n",
        "            'routings': self.routings\n",
        "        }\n",
        "        base_config = super(CapsuleLayer, self).get_config()\n",
        "        return dict(list(base_config.items()) + list(config.items()))\n",
        "\n",
        "\n",
        "def PrimaryCap(inputs, dim_capsule, n_channels, kernel_size, strides, padding):\n",
        "    \"\"\"\n",
        "    Apply Conv2D `n_channels` times and concatenate all capsules\n",
        "    :param inputs: 4D tensor, shape=[None, width, height, channels]\n",
        "    :param dim_capsule: the dim of the output vector of capsule\n",
        "    :param n_channels: the number of types of capsules\n",
        "    :return: output tensor, shape=[None, num_capsule, dim_capsule]\n",
        "    \"\"\"\n",
        "    output = layers.Conv2D(filters=dim_capsule*n_channels, kernel_size=kernel_size, strides=strides, padding=padding,\n",
        "                           name='primarycap_conv2d')(inputs)\n",
        "    outputs = layers.Reshape(target_shape=[-1, dim_capsule], name='primarycap_reshape')(output)\n",
        "    return layers.Lambda(squash, name='primarycap_squash')(outputs)\n",
        "\n",
        "\n",
        "\"\"\"\n",
        "# The following is another way to implement primary capsule layer. This is much slower.\n",
        "# Apply Conv2D `n_channels` times and concatenate all capsules\n",
        "def PrimaryCap(inputs, dim_capsule, n_channels, kernel_size, strides, padding):\n",
        "    outputs = []\n",
        "    for _ in range(n_channels):\n",
        "        output = layers.Conv2D(filters=dim_capsule, kernel_size=kernel_size, strides=strides, padding=padding)(inputs)\n",
        "        outputs.append(layers.Reshape([output.get_shape().as_list()[1] ** 2, dim_capsule])(output))\n",
        "    outputs = layers.Concatenate(axis=1)(outputs)\n",
        "    return layers.Lambda(squash)(outputs)\n",
        "\"\"\""
      ],
      "execution_count": 93,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.2.0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'\\n# The following is another way to implement primary capsule layer. This is much slower.\\n# Apply Conv2D `n_channels` times and concatenate all capsules\\ndef PrimaryCap(inputs, dim_capsule, n_channels, kernel_size, strides, padding):\\n    outputs = []\\n    for _ in range(n_channels):\\n        output = layers.Conv2D(filters=dim_capsule, kernel_size=kernel_size, strides=strides, padding=padding)(inputs)\\n        outputs.append(layers.Reshape([output.get_shape().as_list()[1] ** 2, dim_capsule])(output))\\n    outputs = layers.Concatenate(axis=1)(outputs)\\n    return layers.Lambda(squash)(outputs)\\n'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 93
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "_cell_guid": "7cd17730-22b6-4ac3-a612-31f18902fa78",
        "_uuid": "61c38c7ee701bb3ee2190263cf907fcdbe40dca2",
        "id": "UDdQcGoko8x8",
        "colab_type": "text"
      },
      "source": [
        "# Build the Model\n",
        "Here we use the layers to build up the model. The model is a bit different from a standard $X\\rightarrow y$  model, it is $(X,y)\\rightarrow (y,X)$ meaning it attempts to predict the class from the image, and then at the same time, using the same capsule reconstruct the image from the class. The approach appears very cGAN-like where the task of reconstructing better helps the model 'understand' the image data better."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "bc101123-d53c-4c2e-a187-101c434885da",
        "_uuid": "2497453eb1895f624ad84617dd98c230f5640304",
        "id": "HCMIciV0o8x9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tensorflow.keras import layers, models, optimizers\n",
        "from tensorflow.keras import backend as K\n",
        "from tensorflow.keras.utils import to_categorical\n",
        "\n",
        "K.set_image_data_format('channels_last')\n",
        "\n",
        "def CapsNet(input_shape, n_class, routings, batch):\n",
        "    \"\"\"\n",
        "    A Capsule Network on MNIST.\n",
        "    :param input_shape: data shape, 3d, [width, height, channels]\n",
        "    :param n_class: number of classes\n",
        "    :param routings: number of routing iterations\n",
        "    :param batch_size: size of batch\n",
        "    :return: Two Keras Models, the first one used for training, and the second one for evaluation.\n",
        "            `eval_model` can also be used for training.\n",
        "    \"\"\"\n",
        "    x = layers.Input(shape=input_shape, batch_size=batch)\n",
        "\n",
        "    # Layer 1: Just a conventional Conv2D layer\n",
        "    conv1 = layers.Conv2D(filters=256, kernel_size=9, strides=1, padding='valid', activation='relu', name='conv1')(x)\n",
        "\n",
        "    # Layer 2: Conv2D layer with `squash` activation, then reshape to [None, num_capsule, dim_capsule]\n",
        "    primarycaps = PrimaryCap(conv1, dim_capsule=8, n_channels=32, kernel_size=9, strides=2, padding='valid')\n",
        "\n",
        "    # Layer 3: Capsule layer. Routing algorithm works here.\n",
        "    digitcaps = CapsuleLayer(num_capsule=n_class, dim_capsule=16, routings=routings, name='digitcaps')(primarycaps)\n",
        "\n",
        "    # Layer 4: This is an auxiliary layer to replace each capsule with its length. Just to match the true label's shape.\n",
        "    # If using tensorflow, this will not be necessary. :)\n",
        "    out_caps = Length(name='capsnet')(digitcaps)\n",
        "\n",
        "    # Decoder network.\n",
        "    y = layers.Input(shape=(n_class,))\n",
        "    masked_by_y = Mask()([digitcaps, y])  # The true label is used to mask the output of capsule layer. For training\n",
        "    masked = Mask()(digitcaps)  # Mask using the capsule with maximal length. For prediction\n",
        "\n",
        "    # Shared Decoder model in training and prediction\n",
        "    decoder = models.Sequential(name='decoder')\n",
        "    decoder.add(layers.Dense(512, activation='relu', input_dim=16 * n_class))\n",
        "    decoder.add(layers.Dense(1024, activation='relu'))\n",
        "    decoder.add(layers.Dense(np.prod(input_shape), activation='sigmoid'))\n",
        "    decoder.add(layers.Reshape(target_shape=input_shape, name='out_recon'))\n",
        "\n",
        "    # Models for training and evaluation (prediction)\n",
        "    train_model = models.Model([x, y], [out_caps, decoder(masked_by_y)])\n",
        "    eval_model = models.Model(x, [out_caps, decoder(masked)])\n",
        "\n",
        "    # manipulate model\n",
        "    noise = layers.Input(shape=(n_class, 16))\n",
        "    noised_digitcaps = layers.Add()([digitcaps, noise])\n",
        "    masked_noised_y = Mask()([noised_digitcaps, y])\n",
        "    manipulate_model = models.Model([x, y, noise], decoder(masked_noised_y))\n",
        "    return train_model, eval_model, manipulate_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "c6d84e5a-c33c-40c8-89c3-aba3454f7025",
        "_uuid": "9f27c6b0623ebffb6c8a24579f9dd4e321d6b1c2",
        "id": "LS09Ic7qo8x_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def margin_loss(y_true, y_pred):\n",
        "    \"\"\"\n",
        "    Margin loss for Eq.(4). When y_true[i, :] contains not just one `1`, this loss should work too. Not test it.\n",
        "    :param y_true: [None, n_classes]\n",
        "    :param y_pred: [None, num_capsule]\n",
        "    :return: a scalar loss value.\n",
        "    \"\"\"\n",
        "    # return tf.reduce_mean(tf.square(y_pred))\n",
        "    L = y_true * tf.square(tf.maximum(0., 0.9 - y_pred)) + \\\n",
        "        0.5 * (1 - y_true) * tf.square(tf.maximum(0., y_pred - 0.1))\n",
        "\n",
        "    return tf.reduce_mean(tf.reduce_sum(L, 1))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "6f168849-f8f9-4241-9ba8-59abc59573f1",
        "_uuid": "d21637e677fca5be415b2ff2e8a94ef7db2ea8a1",
        "id": "oa3i75Dao8yC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# define model\n",
        "#(input_shape, n_class, routings, batch_size):\n",
        "#model = CapsNet(input_shape=[28, 28, 1],\n",
        "#                n_class=10,\n",
        "#                routings=3,batch_size=100)\n",
        "#model.summary()\n",
        "#try:\n",
        "#    plot_model(model, to_file='model.png', show_shapes=True)\n",
        "#except Exception as e:\n",
        "#    print('No fancy plot {}'.format(e))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "_cell_guid": "7d75c4e0-ffca-45fc-8acc-620fe2825f82",
        "_uuid": "3e810fabad89045b7f4b51fe8540164f90e7698c",
        "id": "LjrtUDsVo8yE",
        "colab_type": "text"
      },
      "source": [
        "# Load MNIST Data\n",
        "Here we load and reformat the Kaggle contest data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "9dbcd67c-8f99-4d8b-8cf7-480d8dc069a4",
        "_uuid": "526436cc40013621251285812ba95725d4a6d749",
        "id": "7UfmktlEo8yF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_mnist():\n",
        "    # the data, shuffled and split between train and test sets\n",
        "    from tensorflow.keras.datasets import mnist\n",
        "    (x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
        "\n",
        "    x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
        "    x_test = x_test.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
        "    y_train = to_categorical(y_train.astype('float32'))\n",
        "    y_test = to_categorical(y_test.astype('float32'))\n",
        "    return (x_train, y_train), (x_test, y_test)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "d6e78705-f643-40c7-9f1a-d0564d64c718",
        "_uuid": "642af6625a73b7e67193b8d96d23c22d30073643",
        "id": "r7bJRPVVo8yI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#x_train = x_train.values.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
        "##x_test = x_test.values.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
        "#y_train = to_categorical(y_train.astype('float32'))\n",
        "#y_test = to_categorical(y_test.astype('float32'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "e698ab37-4e74-43e6-b35b-a0507449d916",
        "_uuid": "c0374dabdf452026e3f231ffe689b5dd9f99288b",
        "id": "TZJyVmido8yK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(model,  # type: models.Model\n",
        "          data, args):\n",
        "    \"\"\"\n",
        "    Training a CapsuleNet\n",
        "    :param model: the CapsuleNet model\n",
        "    :param data: a tuple containing training and testing data, like `((x_train, y_train), (x_test, y_test))`\n",
        "    :param args: arguments\n",
        "    :return: The trained model\n",
        "    \"\"\"\n",
        "    # unpacking the data\n",
        "    (x_train, y_train), (x_test, y_test) = data\n",
        "\n",
        "    # callbacks\n",
        "    log = callbacks.CSVLogger(args.save_dir + '/log.csv')\n",
        "    checkpoint = callbacks.ModelCheckpoint(args.save_dir + '/weights-{epoch:02d}.h5', monitor='val_capsnet_acc',\n",
        "                                           save_best_only=True, save_weights_only=True, verbose=1)\n",
        "    lr_decay = callbacks.LearningRateScheduler(schedule=lambda epoch: args.lr * (args.lr_decay ** epoch))\n",
        "\n",
        "    # compile the model\n",
        "    model.compile(optimizer=optimizers.Adam(lr=args.lr),\n",
        "                  loss=[margin_loss, 'mse'],\n",
        "                  loss_weights=[1., args.lam_recon],\n",
        "                  metrics={'capsnet': 'accuracy'})\n",
        "\n",
        "    \"\"\"\n",
        "    # Training without data augmentation:\n",
        "    model.fit([x_train, y_train], [y_train, x_train], batch_size=args.batch_size, epochs=args.epochs,\n",
        "              validation_data=[[x_test, y_test], [y_test, x_test]], callbacks=[log, tb, checkpoint, lr_decay])\n",
        "    \"\"\"\n",
        "\n",
        "    # Begin: Training with data augmentation ---------------------------------------------------------------------#\n",
        "    def train_generator(x, y, batch_size, shift_fraction=0.):\n",
        "        train_datagen = ImageDataGenerator(width_shift_range=shift_fraction,\n",
        "                                           height_shift_range=shift_fraction)  # shift up to 2 pixel for MNIST\n",
        "        generator = train_datagen.flow(x, y, batch_size=batch_size)\n",
        "        while 1:\n",
        "            x_batch, y_batch = generator.next()\n",
        "            yield (x_batch, y_batch), (y_batch, x_batch)\n",
        "\n",
        "    # Training with data augmentation. If shift_fraction=0., no augmentation.\n",
        "    model.fit(train_generator(x_train, y_train, args.batch_size, args.shift_fraction),\n",
        "              steps_per_epoch=int(y_train.shape[0] / args.batch_size),\n",
        "              epochs=args.epochs,\n",
        "              validation_data=((x_test, y_test), (y_test, x_test)), batch_size=args.batch_size,\n",
        "              callbacks=[log, checkpoint, lr_decay])\n",
        "    # End: Training with data augmentation -----------------------------------------------------------------------#\n",
        "\n",
        "    model.save_weights(args.save_dir + '/trained_model.h5')\n",
        "    print('Trained model saved to \\'%s/trained_model.h5\\'' % args.save_dir)\n",
        "\n",
        "    from utils import plot_log\n",
        "    plot_log(args.save_dir + '/log.csv', show=True)\n",
        "\n",
        "    return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l-i3c_UESE2P",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 510
        },
        "outputId": "6d77ccfb-a44d-4802-c6c7-9eed6c98753b"
      },
      "source": [
        " # load data\n",
        "(x_train, y_train), (x_test, y_test) = load_mnist()\n",
        "\n",
        "    # define model\n",
        "model, eval_model, manipulate_model = CapsNet(input_shape=x_train.shape[1:],\n",
        "                                                  n_class=len(np.unique(np.argmax(y_train, 1))),\n",
        "                                                  routings=3,\n",
        "                                                  batch=100)\n",
        "model.summary()"
      ],
      "execution_count": 97,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"model\"\n",
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            "input_2 (InputLayer)            [(100, 28, 28, 1)]   0                                            \n",
            "__________________________________________________________________________________________________\n",
            "conv1 (Conv2D)                  (100, 20, 20, 256)   20992       input_2[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "primarycap_conv2d (Conv2D)      (100, 6, 6, 256)     5308672     conv1[0][0]                      \n",
            "__________________________________________________________________________________________________\n",
            "primarycap_reshape (Reshape)    (100, 1152, 8)       0           primarycap_conv2d[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "primarycap_squash (Lambda)      (100, 1152, 8)       0           primarycap_reshape[0][0]         \n",
            "__________________________________________________________________________________________________\n",
            "digitcaps (CapsuleLayer)        (100, 10, 16)        1474560     primarycap_squash[0][0]          \n",
            "__________________________________________________________________________________________________\n",
            "input_3 (InputLayer)            [(None, 10)]         0                                            \n",
            "__________________________________________________________________________________________________\n",
            "mask (Mask)                     (100, 160)           0           digitcaps[0][0]                  \n",
            "                                                                 input_3[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "capsnet (Length)                (100, 10)            0           digitcaps[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "decoder (Sequential)            multiple             1411344     mask[0][0]                       \n",
            "==================================================================================================\n",
            "Total params: 8,215,568\n",
            "Trainable params: 8,215,568\n",
            "Non-trainable params: 0\n",
            "__________________________________________________________________________________________________\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "2fda2834-7c9e-4ed6-b322-0b9e86415201",
        "_uuid": "07bbb33aaa1c7ec875798359e300bbcdba374659",
        "id": "8NCOZGHlo8yM",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 197
        },
        "outputId": "ee3ec00a-261d-479b-e851-9342766a8b92"
      },
      "source": [
        "#train(model=model, data=((x_train, y_train), (x_test[:60], y_test[:60])), \n",
        " #     epoch_size_frac = 0.5) # do 10% of an epoch (takes too long)\n",
        "train(model=model, data=((x_train, y_train), (x_test, y_test)), args=args)"
      ],
      "execution_count": 99,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-99-cdb4c7a47a48>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#train(model=model, data=((x_train, y_train), (x_test[:60], y_test[:60])),\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m  \u001b[0;31m#     epoch_size_frac = 0.5) # do 10% of an epoch (takes too long)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_train\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx_test\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_test\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m: name 'args' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "169b7f78-12c7-4fed-886e-60024fe59339",
        "_uuid": "02b7db879a533e7bfb3116522bebf3867b23498c",
        "id": "0zSuXAxxo8yQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def combine_images(generated_images):\n",
        "    num = generated_images.shape[0]\n",
        "    width = int(np.sqrt(num))\n",
        "    height = int(np.ceil(float(num)/width))\n",
        "    shape = generated_images.shape[1:3]\n",
        "    image = np.zeros((height*shape[0], width*shape[1]),\n",
        "                     dtype=generated_images.dtype)\n",
        "    for index, img in enumerate(generated_images):\n",
        "        i = int(index/width)\n",
        "        j = index % width\n",
        "        image[i*shape[0]:(i+1)*shape[0], j*shape[1]:(j+1)*shape[1]] = \\\n",
        "            img[:, :, 0]\n",
        "    return image\n",
        "\n",
        "def test(model, data, args):\n",
        "    x_test, y_test = data\n",
        "    y_pred, x_recon = model.predict(x_test, batch_size=100)\n",
        "    print('-' * 30 + 'Begin: test' + '-' * 30)\n",
        "    print('Test acc:', np.sum(np.argmax(y_pred, 1) == np.argmax(y_test, 1)) / y_test.shape[0])\n",
        "\n",
        "    img = combine_images(np.concatenate([x_test[:50], x_recon[:50]]))\n",
        "    image = img * 255\n",
        "    Image.fromarray(image.astype(np.uint8)).save(args.save_dir + \"/real_and_recon.png\")\n",
        "    print()\n",
        "    print('Reconstructed images are saved to %s/real_and_recon.png' % args.save_dir)\n",
        "    print('-' * 30 + 'End: test' + '-' * 30)\n",
        "    plt.imshow(plt.imread(args.save_dir + \"/real_and_recon.png\"))\n",
        "    plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "_cell_guid": "2706080d-2d50-4876-bd4d-c0f3c4ce87b5",
        "_uuid": "9afcde53f3cf3ba3eeeed1f083241874b4cd84e2",
        "id": "lMRMqKA9o8yS",
        "colab_type": "text"
      },
      "source": [
        "# Show the results on the hold-out\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "56a6fe58-fe97-45c8-95da-223297842f79",
        "_uuid": "25ac7feb2d111b82e755169288ffd47ebc4d196a",
        "id": "PL8it8clo8yT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test(model=eval_model, data=(x_test, y_test), args=args)\n",
        "#test(model=model, data=(x_test[:100], y_test[:100]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "_cell_guid": "bf642928-3097-4b69-86b4-a2d1f02c1c38",
        "_uuid": "0c1bea7f648c238828af6ceb1ccd3f5fd6d1f075",
        "id": "T3HBtwLJo8yV",
        "colab_type": "text"
      },
      "source": [
        "# Apply Model to the Competition Data\n",
        "Here we apply the model to the compitition data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "e3777e2b-6aff-45e2-b1dd-bee2f0b8eadd",
        "_uuid": "532948f338ea8cc3868ec0511f8866583fa1ab84",
        "id": "ql-25FZEo8yV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data_test = pd.read_csv('../input/test.csv')\n",
        "data_test = data_test.values.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
        "y_pred, _ = model.predict([data_test, \n",
        "                           np.zeros((data_test.shape[0],10))], # empty values for the second vector \n",
        "                           batch_size = 32, verbose = True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "_cell_guid": "0d4794e0-f351-4ad9-beb9-7e76ff457e0a",
        "_uuid": "c00b9d785f9bd7bef37b4f96bc137d18c1219114",
        "id": "X5nhkMgHo8yY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with open('submission.csv', 'w') as out_file:\n",
        "    out_file.write('ImageId,Label\\n')\n",
        "    for img_id, guess_label in enumerate(np.argmax(y_pred,1),1):\n",
        "        out_file.write('%d,%d\\n' % (img_id, guess_label))"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}