Back to Blog

Running CapsuleNet on TensorFlow

November 27, 2017
|
Written by
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

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.

Let’s import all the modules.

Initialise the Capsule Network

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

Training the Capsule Network

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.

Training Process

Training process view

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

For any queries , reach out to me via LinkedIn , Twitter or email me on rahul@botsupply.ai, will love to hear from you. 

Get 30 minutes Free expert advice. No cost, no obligation and no sales pitch.

Schedule a One-on-One