Running CapsuleNet on TensorFlow

Geoffrey Hinton in his lecture

So now we all know that Capsule Networks (by Geoffrey Hinton) is shaking up the AI space and literature states that it will push the limits of Convolutional Neural Network (CNN) to the next level. There are lot of Medium posts, articles and research papers available that discuss about the theory and how it is better than traditional CNN’s.
So I am not going to cover that part, instead I would try to implement the CpNet on TensorFlow using Google’s amazing internal tool called Colaboratory.

Few links you can follow to understand the theory part of CpNet:

• Geoffrey Hinton talk “What is wrong with convolutional neural nets ?”
• Capsule Networks Are Shaking up AI
• Dynamic Routing Between Capsules

Let’s code the network.

Before starting, you can follow my CoLab Notebook and execute the following code.

CoLab URL : https://goo.gl/43Jvju

Now we will clone the repository and install the dependencies. Then we will take the MNIST dataset from the repository and move it out to the parent directory.

!git clone https://github.com/bourdakos1/capsule-networks.git
!pip install -r capsule-networks/requirements.txt
!touch capsule-networks/__init__.py
!mv capsule-networks capsule
!mv capsule/data/ ./data/
! ls

Let’s import all the modules.

import os
import tensorflow as tf
from tqdm import tqdm
from capsule.config import cfg
from capsule.utils import load_mnist
from capsule.capsNet import CapsNet

Initialise the Capsule Network

capsNet = CapsNet(is_training=cfg.is_training)

This is how Capsule Network (CpNet) looks like on Tensorboard graph.

Training the Capsule Network

tf.logging.info('Graph loaded')
sv = tf.train.Supervisor(graph=capsNet.graph,
                             logdir=cfg.logdir,
                             save_model_secs=0)
path = cfg.results + '/accuracy.csv'
if not os.path.exists(cfg.results):
    os.mkdir(cfg.results)
elif os.path.exists(path):
    os.remove(path)
fd_results = open(path, 'w')
fd_results.write('step,test_acc\n')

And creating TF session and running the epochs.
By default the model will be trained for 50 epochs at a batch size of 128. You are always welcome to change the config and try new combinations of hyper parameters.

with sv.managed_session() as sess:
    num_batch = int(60000 / cfg.batch_size)
    num_test_batch = 10000 // cfg.batch_size
    teX, teY = load_mnist(cfg.dataset, False)
    for epoch in range(cfg.epoch):
        if sv.should_stop():
            break
        for step in tqdm(range(num_batch), total=num_batch, ncols=70, leave=False, unit='b'):
            global_step = sess.run(capsNet.global_step)
            sess.run(capsNet.train_op)
if step % cfg.train_sum_freq == 0:
                _, summary_str = sess.run([capsNet.train_op, capsNet.train_summary])
                sv.summary_writer.add_summary(summary_str, global_step)
if (global_step + 1) % cfg.test_sum_freq == 0:
                test_acc = 0
                for i in range(num_test_batch):
                    start = i * cfg.batch_size
                    end = start + cfg.batch_size
                    test_acc += sess.run(capsNet.batch_accuracy, {capsNet.X: teX[start:end], capsNet.labels: teY[start:end]})
                test_acc = test_acc / (cfg.batch_size * num_test_batch)
                fd_results.write(str(global_step + 1) + ',' + str(test_acc) + '\n')
                fd_results.flush()
if epoch % cfg.save_freq == 0:
            sv.saver.save(sess, cfg.logdir + '/model_epoch_%04d_step_%02d' % (epoch, global_step))
fd_results.close()
tf.logging.info('Training done')

Training process

For running 50 epochs of Capsule network it took me around 6 hours on NVIDIA TitanXp card.
But after successful training I was able to achieve 0.0038874 total loss which is incredible. 😃 💯

Total loss plot

Download my trained Capsule model

CpNet Model URL : https://goo.gl/DN7SS3

YOU NEVER KNOW UNTIL YOU TRY IT

If you liked this article, my notebook and models, please applause👏 and share with others!

For any queries , reach out to me via LinkedIn , Twitter or email me on rahul@botsupply.ai.