{"cells":[{"metadata":{"_cell_guid":"7f4c4917-22fb-4b10-879e-51c600e6c3af","_uuid":"9b086e5ec535ad75eada3ca72bf5e6534251074f"},"cell_type":"markdown","source":"# Overview\n\nThe 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\nResult:\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"},{"metadata":{"_cell_guid":"e30e2e10-a909-485d-be9d-dc6f592911a7","_uuid":"c7e569699c6d067cd9fdf9c77299775e399b2ef3","trusted":true},"cell_type":"code","source":"import numpy as np\nimport os\nimport pandas as pd\nfrom keras.preprocessing.image import ImageDataGenerator\nfrom keras import callbacks\nfrom keras.utils.vis_utils import plot_model","execution_count":1,"outputs":[{"output_type":"stream","text":"Using TensorFlow backend.\n","name":"stderr"}]},{"metadata":{"_cell_guid":"9776fd57-44e0-4211-a7a5-c7e647a10704","_uuid":"f4b5499a472b312d5c5f0274ad429567aced6841"},"cell_type":"markdown","source":"# Capsule Layers \nHere is the implementation of the necessary layers for the CapsuleNet. These are not optimized yet and can be made significantly more performant. "},{"metadata":{"_cell_guid":"167d39ca-ee32-4eec-a83b-86194252b14f","collapsed":true,"_uuid":"90c180a9a8c20e3fb8a93c3eb42588927cfcd6b6","trusted":false},"cell_type":"code","source":"import keras.backend as K\nimport tensorflow as tf\nfrom keras import initializers, layers\n\nclass 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    inputs: shape=[dim_1, ..., dim_{n-1}, dim_n]\n    output: shape=[dim_1, ..., dim_{n-1}]\n    \"\"\"\n    def call(self, inputs, **kwargs):\n        return K.sqrt(K.sum(K.square(inputs), -1))\n\n    def compute_output_shape(self, input_shape):\n        return input_shape[:-1]\n\nclass Mask(layers.Layer):\n    \"\"\"\n    Mask a Tensor with shape=[None, d1, d2] by the max value in axis=1.\n    Output shape: [None, d2]\n    \"\"\"\n    def call(self, inputs, **kwargs):\n        # use true label to select target capsule, shape=[batch_size, num_capsule]\n        if type(inputs) is list:  # true label is provided with shape = [batch_size, 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 vectors of capsules\n            x = inputs\n            # Enlarge the range of values in x to make max(new_x)=1 and others < 0\n            x = (x - K.max(x, 1, True)) / K.epsilon() + 1\n            mask = K.clip(x, 0, 1)  # the max value in x clipped to 1 and other to 0\n\n        # masked inputs, shape = [batch_size, dim_vector]\n        inputs_masked = K.batch_dot(inputs, mask, [1, 1])\n        return inputs_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]])\n        else:\n            return tuple([None, input_shape[-1]])\n\n\ndef 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 = K.sum(K.square(vectors), axis, keepdims=True)\n    scale = s_squared_norm / (1 + s_squared_norm) / K.sqrt(s_squared_norm)\n    return scale * vectors\n\n\nclass 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_vector] and output shape = \\\n    [None, num_capsule, dim_vector]. For Dense Layer, input_dim_vector = dim_vector = 1.\n    \n    :param num_capsule: number of capsules in this layer\n    :param dim_vector: dimension of the output vectors of the capsules in this layer\n    :param num_routings: number of iterations for the routing algorithm\n    \"\"\"\n    def __init__(self, num_capsule, dim_vector, num_routing=3,\n                 kernel_initializer='glorot_uniform',\n                 bias_initializer='zeros',\n                 **kwargs):\n        super(CapsuleLayer, self).__init__(**kwargs)\n        self.num_capsule = num_capsule\n        self.dim_vector = dim_vector\n        self.num_routing = num_routing\n        self.kernel_initializer = initializers.get(kernel_initializer)\n        self.bias_initializer = initializers.get(bias_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_vector]\"\n        self.input_num_capsule = input_shape[1]\n        self.input_dim_vector = input_shape[2]\n\n        # Transform matrix\n        self.W = self.add_weight(shape=[self.input_num_capsule, self.num_capsule, self.input_dim_vector, self.dim_vector],\n                                 initializer=self.kernel_initializer,\n                                 name='W')\n\n        # Coupling coefficient. The redundant dimensions are just to facilitate subsequent matrix calculation.\n        self.bias = self.add_weight(shape=[1, self.input_num_capsule, self.num_capsule, 1, 1],\n                                    initializer=self.bias_initializer,\n                                    name='bias',\n                                    trainable=False)\n        self.built = True\n\n    def call(self, inputs, training=None):\n        # inputs.shape=[None, input_num_capsule, input_dim_vector]\n        # Expand dims to [None, input_num_capsule, 1, 1, input_dim_vector]\n        inputs_expand = K.expand_dims(K.expand_dims(inputs, 2), 2)\n\n        # Replicate num_capsule dimension to prepare being multiplied by W\n        # Now it has shape = [None, input_num_capsule, num_capsule, 1, input_dim_vector]\n        inputs_tiled = K.tile(inputs_expand, [1, 1, self.num_capsule, 1, 1])\n\n        \"\"\"  \n        # Compute `inputs * W` by expanding the first dim of W. More time-consuming and need batch_size.\n        # Now W has shape  = [batch_size, input_num_capsule, num_capsule, input_dim_vector, dim_vector]\n        w_tiled = K.tile(K.expand_dims(self.W, 0), [self.batch_size, 1, 1, 1, 1])\n        \n        # Transformed vectors, inputs_hat.shape = [None, input_num_capsule, num_capsule, 1, dim_vector]\n        inputs_hat = K.batch_dot(inputs_tiled, w_tiled, [4, 3])\n        \"\"\"\n        # Compute `inputs * W` by scanning inputs_tiled on dimension 0. This is faster but requires Tensorflow.\n        # inputs_hat.shape = [None, input_num_capsule, num_capsule, 1, dim_vector]\n        inputs_hat = tf.scan(lambda ac, x: K.batch_dot(x, self.W, [3, 2]),\n                             elems=inputs_tiled,\n                             initializer=K.zeros([self.input_num_capsule, self.num_capsule, 1, self.dim_vector]))\n        \"\"\"\n        # Routing algorithm V1. Use tf.while_loop in a dynamic way.\n        def body(i, b, outputs):\n            c = tf.nn.softmax(self.bias, dim=2)  # dim=2 is the num_capsule dimension\n            outputs = squash(K.sum(c * inputs_hat, 1, keepdims=True))\n            b = b + K.sum(inputs_hat * outputs, -1, keepdims=True)\n            return [i-1, b, outputs]\n\n        cond = lambda i, b, inputs_hat: i > 0\n        loop_vars = [K.constant(self.num_routing), self.bias, K.sum(inputs_hat, 1, keepdims=True)]\n        _, _, outputs = tf.while_loop(cond, body, loop_vars)\n        \"\"\"\n        # Routing algorithm V2. Use iteration. V2 and V1 both work without much difference on performance\n        assert self.num_routing > 0, 'The num_routing should be > 0.'\n        for i in range(self.num_routing):\n            c = tf.nn.softmax(self.bias, dim=2)  # dim=2 is the num_capsule dimension\n            # outputs.shape=[None, 1, num_capsule, 1, dim_vector]\n            outputs = squash(K.sum(c * inputs_hat, 1, keepdims=True))\n\n            # last iteration needs not compute bias which will not be passed to the graph any more anyway.\n            if i != self.num_routing - 1:\n                # self.bias = K.update_add(self.bias, K.sum(inputs_hat * outputs, [0, -1], keepdims=True))\n                self.bias += K.sum(inputs_hat * outputs, -1, keepdims=True)\n            # tf.summary.histogram('BigBee', self.bias)  # for debugging\n        return K.reshape(outputs, [-1, self.num_capsule, self.dim_vector])\n\n    def compute_output_shape(self, input_shape):\n        return tuple([None, self.num_capsule, self.dim_vector])\n\n\ndef PrimaryCap(inputs, dim_vector, 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_vector: 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_vector]\n    \"\"\"\n    output = layers.Conv2D(filters=dim_vector*n_channels, kernel_size=kernel_size, strides=strides, padding=padding)(inputs)\n    outputs = layers.Reshape(target_shape=[-1, dim_vector])(output)\n    return layers.Lambda(squash)(outputs)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"7cd17730-22b6-4ac3-a612-31f18902fa78","_uuid":"61c38c7ee701bb3ee2190263cf907fcdbe40dca2"},"cell_type":"markdown","source":"# Build the Model\nHere 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."},{"metadata":{"_cell_guid":"bc101123-d53c-4c2e-a187-101c434885da","collapsed":true,"_uuid":"2497453eb1895f624ad84617dd98c230f5640304","trusted":false},"cell_type":"code","source":"from keras import layers, models\nfrom keras import backend as K\nfrom keras.utils import to_categorical\ndef CapsNet(input_shape, n_class, num_routing):\n    \"\"\"\n    A Capsule Network on MNIST.\n    :param input_shape: data shape, 4d, [None, width, height, channels]\n    :param n_class: number of classes\n    :param num_routing: number of routing iterations\n    :return: A Keras Model with 2 inputs and 2 outputs\n    \"\"\"\n    x = layers.Input(shape=input_shape)\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_vector]\n    primarycaps = PrimaryCap(conv1, dim_vector=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_vector=16, num_routing=num_routing, 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='out_caps')(digitcaps)\n\n    # Decoder network.\n    y = layers.Input(shape=(n_class,))\n    masked = Mask()([digitcaps, y])  # The true label is used to mask the output of capsule layer.\n    x_recon = layers.Dense(512, activation='relu')(masked)\n    x_recon = layers.Dense(1024, activation='relu')(x_recon)\n    x_recon = layers.Dense(784, activation='sigmoid')(x_recon)\n    x_recon = layers.Reshape(target_shape=[28, 28, 1], name='out_recon')(x_recon)\n\n    # two-input-two-output keras Model\n    return models.Model([x, y], [out_caps, x_recon])","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"c6d84e5a-c33c-40c8-89c3-aba3454f7025","collapsed":true,"_uuid":"9f27c6b0623ebffb6c8a24579f9dd4e321d6b1c2","trusted":false},"cell_type":"code","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    L = y_true * K.square(K.maximum(0., 0.9 - y_pred)) + \\\n        0.5 * (1 - y_true) * K.square(K.maximum(0., y_pred - 0.1))\n\n    return K.mean(K.sum(L, 1))","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"6f168849-f8f9-4241-9ba8-59abc59573f1","collapsed":true,"_uuid":"d21637e677fca5be415b2ff2e8a94ef7db2ea8a1","trusted":false},"cell_type":"code","source":"# define model\nmodel = CapsNet(input_shape=[28, 28, 1],\n                n_class=10,\n                num_routing=3)\nmodel.summary()\ntry:\n    plot_model(model, to_file='model.png', show_shapes=True)\nexcept Exception as e:\n    print('No fancy plot {}'.format(e))","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"7d75c4e0-ffca-45fc-8acc-620fe2825f82","_uuid":"3e810fabad89045b7f4b51fe8540164f90e7698c"},"cell_type":"markdown","source":"# Load MNIST Data\nHere we load and reformat the Kaggle contest data"},{"metadata":{"_cell_guid":"9dbcd67c-8f99-4d8b-8cf7-480d8dc069a4","collapsed":true,"_uuid":"526436cc40013621251285812ba95725d4a6d749","trusted":false},"cell_type":"code","source":"from sklearn.model_selection import train_test_split\ndata_train = pd.read_csv('../input/train.csv')\nX_full = data_train.iloc[:,1:]\ny_full = data_train.iloc[:,:1]\nx_train, x_test, y_train, y_test = train_test_split(X_full, y_full, test_size = 0.3)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"d6e78705-f643-40c7-9f1a-d0564d64c718","collapsed":true,"_uuid":"642af6625a73b7e67193b8d96d23c22d30073643","trusted":false},"cell_type":"code","source":"x_train = x_train.values.reshape(-1, 28, 28, 1).astype('float32') / 255.\nx_test = x_test.values.reshape(-1, 28, 28, 1).astype('float32') / 255.\ny_train = to_categorical(y_train.astype('float32'))\ny_test = to_categorical(y_test.astype('float32'))","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"e698ab37-4e74-43e6-b35b-a0507449d916","collapsed":true,"_uuid":"c0374dabdf452026e3f231ffe689b5dd9f99288b","trusted":false},"cell_type":"code","source":"def train(model, data, epoch_size_frac=1.0):\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('log.csv')\n    checkpoint = callbacks.ModelCheckpoint('weights-{epoch:02d}.h5',\n                                           save_best_only=True, save_weights_only=True, verbose=1)\n    lr_decay = callbacks.LearningRateScheduler(schedule=lambda epoch: 0.001 * np.exp(-epoch / 10.))\n\n    # compile the model\n    model.compile(optimizer='adam',\n                  loss=[margin_loss, 'mse'],\n                  loss_weights=[1., 0.0005],\n                  metrics={'out_caps': '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])\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., also no augmentation.\n    model.fit_generator(generator=train_generator(x_train, y_train, 64, 0.1),\n                        steps_per_epoch=int(epoch_size_frac*y_train.shape[0] / 64),\n                        epochs=1,\n                        validation_data=[[x_test, y_test], [y_test, x_test]],\n                        callbacks=[log, checkpoint, lr_decay])\n    # -----------------------------------End: Training with data augmentation -----------------------------------#\n\n    model.save_weights('trained_model.h5')\n    print('Trained model saved to \\'trained_model.h5\\'')\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"2fda2834-7c9e-4ed6-b322-0b9e86415201","collapsed":true,"_uuid":"07bbb33aaa1c7ec875798359e300bbcdba374659","trusted":false},"cell_type":"code","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)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"169b7f78-12c7-4fed-886e-60024fe59339","collapsed":true,"_uuid":"02b7db879a533e7bfb3116522bebf3867b23498c","trusted":false},"cell_type":"code","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\ndef test(model, data):\n    x_test, y_test = data\n    y_pred, x_recon = model.predict([x_test, y_test], batch_size=100)\n    print('-'*50)\n    print('Test acc:', np.sum(np.argmax(y_pred, 1) == np.argmax(y_test, 1))/y_test.shape[0])\n\n    import matplotlib.pyplot as plt\n    from PIL import Image\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(\"real_and_recon.png\")\n    print()\n    print('Reconstructed images are saved to ./real_and_recon.png')\n    print('-'*50)\n    plt.imshow(plt.imread(\"real_and_recon.png\", ))\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"2706080d-2d50-4876-bd4d-c0f3c4ce87b5","_uuid":"9afcde53f3cf3ba3eeeed1f083241874b4cd84e2"},"cell_type":"markdown","source":"# Show the results on the hold-out\n"},{"metadata":{"_cell_guid":"56a6fe58-fe97-45c8-95da-223297842f79","collapsed":true,"_uuid":"25ac7feb2d111b82e755169288ffd47ebc4d196a","trusted":false},"cell_type":"code","source":"test(model=model, data=(x_test[:100], y_test[:100]))","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"bf642928-3097-4b69-86b4-a2d1f02c1c38","_uuid":"0c1bea7f648c238828af6ceb1ccd3f5fd6d1f075"},"cell_type":"markdown","source":"# Apply Model to the Competition Data\nHere we apply the model to the compitition data"},{"metadata":{"_cell_guid":"e3777e2b-6aff-45e2-b1dd-bee2f0b8eadd","collapsed":true,"_uuid":"532948f338ea8cc3868ec0511f8866583fa1ab84","trusted":false},"cell_type":"code","source":"data_test = pd.read_csv('../input/test.csv')\ndata_test = data_test.values.reshape(-1, 28, 28, 1).astype('float32') / 255.\ny_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":null,"outputs":[]},{"metadata":{"_cell_guid":"0d4794e0-f351-4ad9-beb9-7e76ff457e0a","_uuid":"c00b9d785f9bd7bef37b4f96bc137d18c1219114","trusted":true},"cell_type":"code","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":null,"outputs":[]}],"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"}},"nbformat":4,"nbformat_minor":1}