capsule-network-dynamic-routing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import keras\n",
    "import keras.backend as K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "(X_train,Y_train),(X_test,Y_test) = keras.datasets.mnist.load_data()\n",
    "X_train_new = np.zeros((X_train.shape[0],32,32))\n",
    "X_test_new = np.zeros((X_test.shape[0],32,32))\n",
    "X_train_new[:,2:-2,2:-2] = X_train\n",
    "X_test_new[:,2:-2,2:-2] = X_test\n",
    "X_train,X_test = X_train_new[...,np.newaxis],X_test_new[...,np.newaxis]\n",
    "X_train,X_test = X_train/255,X_test/255\n",
    "Y_train = keras.utils.to_categorical(Y_train, num_classes=10)\n",
    "Y_test = keras.utils.to_categorical(Y_test, num_classes=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class CapsuleConv2D(keras.layers.Layer):\n",
    "    \n",
    "    def __init__(self, channels, kernel_size, strides=(1,1),\n",
    "                 steps=3, mat_size=(4,4), **kwargs):\n",
    "        self.channels = channels\n",
    "        self.kernel_size = kernel_size\n",
    "        self.strides = strides\n",
    "        self.steps = steps\n",
    "        self.mat_size = mat_size\n",
    "        super().__init__(**kwargs)\n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        self.kernel = self.add_weight(\n",
    "            shape=(self.kernel_size[0]*self.kernel_size[1]*input_shape[3],\n",
    "                   self.channels, *self.mat_size),\n",
    "            initializer=keras.initializers.TruncatedNormal(\n",
    "                mean=0., stddev=1.), name='kernel')\n",
    "        super().build(input_shape)\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        input_shape = inputs.shape\n",
    "        \n",
    "        poses = inputs\n",
    "        input_channel_size = poses.shape[3]\n",
    "        \n",
    "        poses = K.reshape(poses, (\n",
    "            -1,*poses.shape[1:3],poses.shape[3]*poses.shape[4]*poses.shape[5]))\n",
    "        def conv_expand(x):\n",
    "            conv_data_size = x.shape[3]\n",
    "            conv_kernel = np.zeros((*self.kernel_size, conv_data_size,\n",
    "                                    self.kernel_size[0]*self.kernel_size[1]))\n",
    "            for i in range(self.kernel_size[0]):\n",
    "                for j in range(self.kernel_size[1]):\n",
    "                    conv_kernel[i,j,:,i*self.kernel_size[0]+j] = 1\n",
    "            x = K.depthwise_conv2d(x, K.constant(conv_kernel),\n",
    "                                       strides=self.strides, padding='same')\n",
    "            x = K.reshape(x, (-1,*x.shape[1:3],conv_data_size,\n",
    "                                     self.kernel_size[0]*self.kernel_size[1]))\n",
    "            x = K.permute_dimensions(x, (0,1,2,4,3))\n",
    "            return x\n",
    "        poses = conv_expand(poses)\n",
    "        \n",
    "        poses = K.reshape(poses, (-1,*poses.shape[1:3],\n",
    "            self.kernel_size[0]*self.kernel_size[1]*input_channel_size,\n",
    "            1,input_shape[-2],input_shape[-1]))\n",
    "        poses = K.tile(poses, (1,1,1,1,self.channels,1,1))\n",
    "        kernel = K.reshape(self.kernel, (1,1,1,*self.kernel.shape))\n",
    "        kernel = K.tile(kernel, K.concatenate([\n",
    "            K.shape(poses)[:3],K.constant(np.ones((4,)),dtype='int32')]))\n",
    "        votes = kernel @ poses\n",
    "        votes = K.reshape(votes, (\n",
    "            -1,*votes.shape[1:5],self.mat_size[0]*input_shape[-1]))\n",
    "        \n",
    "        b = K.zeros((1,*votes.shape[1:5]))\n",
    "        b = K.tile(b, K.concatenate([\n",
    "            K.shape(votes)[:1],K.constant(np.ones((4,)),dtype='int32')]))\n",
    "        for t in range(self.steps):\n",
    "            c = K.softmax(b, axis=-1)\n",
    "            s = K.expand_dims(c) * votes\n",
    "            s = K.sum(s, axis=3)\n",
    "            s_norm2 = K.sum(K.square(s),axis=-1,keepdims=True)\n",
    "            v = (s_norm2/(1+s_norm2))\n",
    "            v = (s/K.clip(K.sqrt(s_norm2),K.epsilon(),None)) * v\n",
    "            if t == self.steps-1: continue\n",
    "            b = b + K.sum(votes*K.expand_dims(v,axis=3),axis=-1)\n",
    "        \n",
    "        poses_new = K.reshape(v,(\n",
    "            -1,*v.shape[1:4],self.mat_size[0],input_shape[-1]))\n",
    "        return poses_new\n",
    "    \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return (input_shape[0],\n",
    "                input_shape[1]//self.strides[0],\n",
    "                input_shape[2]//self.strides[1],self.channels,\n",
    "                self.mat_size[0],input_shape[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "X = X_input = keras.layers.Input(X_train.shape[1:])\n",
    "X = keras.layers.BatchNormalization()(X)\n",
    "X = keras.layers.Conv2D(32, (3,3), strides=(2,2), padding='same',\n",
    "    activation='relu', kernel_initializer='he_normal')(X)\n",
    "X = keras.layers.Conv2D(32, (3,3), strides=(1,1), padding='same',\n",
    "    activation='relu', kernel_initializer='he_normal')(X)\n",
    "X = keras.layers.Conv2D(64, (3,3), strides=(1,1), padding='same',\n",
    "    activation='relu', kernel_initializer='he_normal')(X)\n",
    "C = keras.Model(X_input, X)\n",
    "C.compile('nadam', 'mse', ['acc'])\n",
    "CF = keras.Model(X_input, X)\n",
    "CF.trainable = False\n",
    "CF.compile('nadam', 'mse', ['acc'])\n",
    "X = X_input = keras.layers.Input(X_train.shape[1:])\n",
    "X = C(X)\n",
    "X = keras.layers.Conv2D(4, (3,3), strides=(2,2), padding='same',\n",
    "    activation='relu', kernel_initializer='he_normal')(X)\n",
    "X = keras.layers.Conv2D(4, (3,3), strides=(2,2), padding='same',\n",
    "    activation='relu', kernel_initializer='he_normal')(X)\n",
    "X = keras.layers.Conv2D(4, (3,3), strides=(2,2), padding='same',\n",
    "    activation='relu', kernel_initializer='he_normal')(X)\n",
    "X = keras.layers.Flatten()(X)\n",
    "X = keras.layers.Dense(Y_train.shape[-1], activation='softmax')(X)\n",
    "CM = keras.Model(X_input, X)\n",
    "CM.compile('nadam', 'categorical_crossentropy', ['acc'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 60000 samples, validate on 10000 samples\n",
      "Epoch 1/5\n",
      "60000/60000 [==============================] - 31s 509us/step - loss: 0.3020 - acc: 0.9074 - val_loss: 0.1276 - val_acc: 0.9608\n",
      "Epoch 2/5\n",
      "60000/60000 [==============================] - 29s 483us/step - loss: 0.1006 - acc: 0.9701 - val_loss: 0.0771 - val_acc: 0.9755\n",
      "Epoch 3/5\n",
      "60000/60000 [==============================] - 29s 482us/step - loss: 0.0748 - acc: 0.9770 - val_loss: 0.0693 - val_acc: 0.9783\n",
      "Epoch 4/5\n",
      "60000/60000 [==============================] - 29s 484us/step - loss: 0.0602 - acc: 0.9823 - val_loss: 0.0601 - val_acc: 0.9816\n",
      "Epoch 5/5\n",
      "60000/60000 [==============================] - 29s 483us/step - loss: 0.0506 - acc: 0.9843 - val_loss: 0.0627 - val_acc: 0.9828\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fe4d45b6a90>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CM.fit(X_train, Y_train,\n",
    "      validation_data=(X_test, Y_test),\n",
    "      batch_size=32, epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_3 (InputLayer)            (None, 32, 32, 1)    0                                            \n",
      "__________________________________________________________________________________________________\n",
      "model_2 (Model)                 (None, 16, 16, 64)   28068       input_3[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_7 (Conv2D)               (None, 16, 16, 32)   2080        model_2[1][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "lambda_1 (Lambda)               (None, 16, 16, 4, 4, 0           conv2d_7[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "capsule_conv2d_1 (CapsuleConv2D (None, 8, 8, 4, 4, 2 2304        lambda_1[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "capsule_conv2d_2 (CapsuleConv2D (None, 4, 4, 4, 4, 2 2304        capsule_conv2d_1[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "capsule_conv2d_3 (CapsuleConv2D (None, 2, 2, 4, 4, 2 2304        capsule_conv2d_2[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "capsule_conv2d_4 (CapsuleConv2D (None, 1, 1, 10, 4,  5760        capsule_conv2d_3[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "lambda_3 (Lambda)               (None, 10)           0           capsule_conv2d_4[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "lambda_2 (Lambda)               (None, 80)           0           capsule_conv2d_4[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "lambda_4 (Lambda)               (None, 10)           0           lambda_3[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_2 (Dense)                 (None, 64)           5184        lambda_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "softmax_1 (Softmax)             (None, 10)           0           lambda_4[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_3 (Dense)                 (None, 10)           650         dense_2[0][0]                    \n",
      "==================================================================================================\n",
      "Total params: 48,654\n",
      "Trainable params: 20,586\n",
      "Non-trainable params: 28,068\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "X = X_input = keras.layers.Input(X_train.shape[1:])\n",
    "X = CF(X)\n",
    "X = keras.layers.Conv2D(4*4*2, (1,1), padding='same', activation='tanh')(X)\n",
    "X = keras.layers.Lambda(\n",
    "    lambda x: K.reshape(x, (-1,*x.shape[1:3],4,4,2)),\n",
    "    output_shape=(int(X.shape[1]),int(X.shape[2]),4,4,2))(X)\n",
    "X = CapsuleConv2D(4, (3,3), strides=(2,2))(X)\n",
    "X = CapsuleConv2D(4, (3,3), strides=(2,2))(X)\n",
    "X = CapsuleConv2D(4, (3,3), strides=(2,2))(X)\n",
    "X = CapsuleConv2D(Y_train.shape[-1], (3,3), strides=(2,2))(X)\n",
    "X_pose = keras.layers.Lambda(\n",
    "    lambda x,d: K.stop_gradient(K.reshape(x,(-1,d*4*2))),\n",
    "    output_shape=(Y_train.shape[-1]*4*2,),\n",
    "    arguments={'d':Y_train.shape[-1]})(X)\n",
    "X_pose = keras.layers.Dense(64, activation='tanh')(X_pose)\n",
    "X_pose = keras.layers.Dense(Y_train.shape[-1], activation='softmax')(X_pose)\n",
    "X = keras.layers.Lambda(\n",
    "    lambda x,d: K.reshape(K.sqrt(K.sum(K.sum(\n",
    "        K.square(x),axis=-1),axis=-1)),(-1,d)),\n",
    "    output_shape=(Y_train.shape[-1],),\n",
    "    arguments={'d':Y_train.shape[-1]})(X)\n",
    "X = keras.layers.Lambda(lambda x: K.print_tensor(x))(X)\n",
    "X = keras.layers.Softmax()(X)\n",
    "M = keras.Model(X_input, [X,X_pose])\n",
    "M_optimizer = keras.optimizers.SGD(momentum=0.9)\n",
    "M.compile(M_optimizer, 'categorical_crossentropy', ['acc'])\n",
    "M.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 60000 samples, validate on 10000 samples\n",
      "Epoch 1/10\n",
      "60000/60000 [==============================] - 889s 15ms/step - loss: 2.0419 - softmax_1_loss: 1.7045 - dense_3_loss: 0.3373 - softmax_1_acc: 0.8431 - dense_3_acc: 0.8950 - val_loss: 1.7459 - val_softmax_1_loss: 1.5886 - val_dense_3_loss: 0.1572 - val_softmax_1_acc: 0.9350 - val_dense_3_acc: 0.9496\n",
      "Epoch 2/10\n",
      "60000/60000 [==============================] - 887s 15ms/step - loss: 1.6862 - softmax_1_loss: 1.5612 - dense_3_loss: 0.1250 - softmax_1_acc: 0.9532 - dense_3_acc: 0.9632 - val_loss: 1.6435 - val_softmax_1_loss: 1.5361 - val_dense_3_loss: 0.1074 - val_softmax_1_acc: 0.9662 - val_dense_3_acc: 0.9698\n",
      "Epoch 3/10\n",
      "60000/60000 [==============================] - 883s 15ms/step - loss: 1.6246 - softmax_1_loss: 1.5271 - dense_3_loss: 0.0975 - softmax_1_acc: 0.9697 - dense_3_acc: 0.9724 - val_loss: 1.6133 - val_softmax_1_loss: 1.5175 - val_dense_3_loss: 0.0957 - val_softmax_1_acc: 0.9723 - val_dense_3_acc: 0.9721\n",
      "Epoch 4/10\n",
      "60000/60000 [==============================] - 886s 15ms/step - loss: 1.5974 - softmax_1_loss: 1.5130 - dense_3_loss: 0.0844 - softmax_1_acc: 0.9736 - dense_3_acc: 0.9755 - val_loss: 1.5919 - val_softmax_1_loss: 1.5084 - val_dense_3_loss: 0.0835 - val_softmax_1_acc: 0.9743 - val_dense_3_acc: 0.9754\n",
      "Epoch 5/10\n",
      "60000/60000 [==============================] - 886s 15ms/step - loss: 1.5793 - softmax_1_loss: 1.5049 - dense_3_loss: 0.0743 - softmax_1_acc: 0.9770 - dense_3_acc: 0.9785 - val_loss: 1.5807 - val_softmax_1_loss: 1.5021 - val_dense_3_loss: 0.0786 - val_softmax_1_acc: 0.9772 - val_dense_3_acc: 0.9789\n",
      "Epoch 6/10\n",
      "60000/60000 [==============================] - 884s 15ms/step - loss: 1.5674 - softmax_1_loss: 1.4994 - dense_3_loss: 0.0680 - softmax_1_acc: 0.9786 - dense_3_acc: 0.9802 - val_loss: 1.5709 - val_softmax_1_loss: 1.4974 - val_dense_3_loss: 0.0735 - val_softmax_1_acc: 0.9799 - val_dense_3_acc: 0.9799\n",
      "Epoch 7/10\n",
      "60000/60000 [==============================] - 887s 15ms/step - loss: 1.5591 - softmax_1_loss: 1.4955 - dense_3_loss: 0.0637 - softmax_1_acc: 0.9815 - dense_3_acc: 0.9820 - val_loss: 1.5750 - val_softmax_1_loss: 1.4962 - val_dense_3_loss: 0.0788 - val_softmax_1_acc: 0.9786 - val_dense_3_acc: 0.9790\n",
      "Epoch 8/10\n",
      "60000/60000 [==============================] - 884s 15ms/step - loss: 1.5512 - softmax_1_loss: 1.4924 - dense_3_loss: 0.0589 - softmax_1_acc: 0.9823 - dense_3_acc: 0.9833 - val_loss: 1.5586 - val_softmax_1_loss: 1.4918 - val_dense_3_loss: 0.0668 - val_softmax_1_acc: 0.9810 - val_dense_3_acc: 0.9822\n",
      "Epoch 9/10\n",
      "60000/60000 [==============================] - 886s 15ms/step - loss: 1.5457 - softmax_1_loss: 1.4899 - dense_3_loss: 0.0559 - softmax_1_acc: 0.9832 - dense_3_acc: 0.9841 - val_loss: 1.5580 - val_softmax_1_loss: 1.4901 - val_dense_3_loss: 0.0679 - val_softmax_1_acc: 0.9814 - val_dense_3_acc: 0.9816\n",
      "Epoch 10/10\n",
      "60000/60000 [==============================] - 884s 15ms/step - loss: 1.5409 - softmax_1_loss: 1.4880 - dense_3_loss: 0.0529 - softmax_1_acc: 0.9838 - dense_3_acc: 0.9846 - val_loss: 1.5533 - val_softmax_1_loss: 1.4887 - val_dense_3_loss: 0.0646 - val_softmax_1_acc: 0.9836 - val_dense_3_acc: 0.9836\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fe44447eba8>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "M.fit(X_train, [Y_train,Y_train],\n",
    "      validation_data=(X_test, [Y_test,Y_test]),\n",
    "      batch_size=8, epochs=10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml",
   "language": "python",
   "name": "ml"
  },
  "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
}