Running CapsuleNet on TensorFlow
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:
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 tqdmfrom 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')
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. 😃 💯
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.