Here is how is how we proceed with the recipe:

1. Clone the code from github under an Apache Licence:

git clone https://github.com/naturomics/CapsNet-Tensorflow.git
 $ cd CapsNet-Tensorflow

2. Download MNIST and create the appropriate structure:

mkdir -p data/mnist
wget -c -P data/mnist \
http://yann.lecun.com/exdb/mnist/{train-images-idx3-ubyte.gz,train-labels-idx1-ubyte.gz,t10k-images-idx3-ubyte.gz,t10k-labels-idx1-ubyte.gz}
gunzip data/mnist/*.gz

3. Start the training process:

python main.py

4. Let's see the code used for defining a capsule. Each capsule takes a 4D tensor as input and returns a 4D tensor. It is possible to define a capsule as either a fully connected network (the DigiCaps) or a convolutional one (the Primary Capsules). Note that Primary is a collection of ConvNets and after them the nonlinear squashing function is applied. Primary capsules will communicate with DigiCaps via dynamic routing:

# capsLayer.py
#
import numpy as np
import tensorflow as tf
from config import cfg
epsilon = 1e-9
class CapsLayer(object):
''' Capsule layer.
Args:
input: A 4-D tensor.
num_outputs: the number of capsule in this layer.
vec_len: integer, the length of the output vector of a capsule.
layer_type: string, one of 'FC' or "CONV", the type of this layer,
fully connected or convolution, for the future expansion capability
with_routing: boolean, this capsule is routing with the
lower-level layer capsule.
Returns:
A 4-D tensor.
'''
def __init__(self, num_outputs, vec_len, with_routing=True, layer_type='FC'):
  self.num_outputs = num_outputs
  self.vec_len = vec_len
  self.with_routing = with_routing
  self.layer_type = layer_type

def __call__(self, input, kernel_size=None, stride=None):
'''
The parameters 'kernel_size' and 'stride' will be used while 'layer_type' equal 'CONV'
'''
  if self.layer_type == 'CONV':
    self.kernel_size = kernel_size
    self.stride = stride

    if not self.with_routing:
    # the PrimaryCaps layer, a convolutional layer
    # input: [batch_size, 20, 20, 256]
      assert input.get_shape() ==  [cfg.batch_size, 20, 20, 256]
      capsules = []
      for i in range(self.vec_len):
        # each capsule i: [batch_size, 6, 6, 32]
        with tf.variable_scope('ConvUnit_' + str(i)):
          caps_i = tf.contrib.layers.conv2d(input,      self.num_outputs,
self.kernel_size, self.stride,
padding="VALID")
          caps_i = tf.reshape(caps_i, shape=(cfg.batch_size, -1, 1, 1))
          capsules.append(caps_i)
      
      assert capsules[0].get_shape() == [cfg.batch_size, 1152, 1, 1]
# [batch_size, 1152, 8, 1]
      capsules = tf.concat(capsules, axis=2)
      capsules = squash(capsules)
      assert capsules.get_shape() == [cfg.batch_size, 1152, 8, 1]
      return(capsules)

  if self.layer_type == 'FC':
    if self.with_routing:
      # the DigitCaps layer, a fully connected layer
      # Reshape the input into [batch_size, 1152, 1, 8, 1]
      self.input = tf.reshape(input, shape=(cfg.batch_size, -1, 1, input.shape[-2].value, 1))
    with tf.variable_scope('routing'):
      # b_IJ: [1, num_caps_l, num_caps_l_plus_1, 1, 1]
      b_IJ = tf.constant(np.zeros([1, input.shape[1].value,    self.num_outputs, 1, 1], dtype=np.float32))
      capsules = routing(self.input, b_IJ)
      capsules = tf.squeeze(capsules, axis=1)
    return(capsules)

5. The routing algorithm is described in the paper Dynamic Routing Between Capsules and the relevant section of the paper has been explained, together with the definition of Equation 2 and Equation 3. The goal of the routing algorithm is to pass information from lower layer capsules into higher-level ones and understand where there is agreement. The agreement is computed by simply using a scalar product between the current output v_j of each capsule j in the layer above, and the prediction  made by the capsule i:

The following method implements the steps described in Procedure 1 in the preceding images. Note that the input is a 4D tensor from 1,152 capsules in the layer l. The output is a Tensor of shape [batch_size, 1, length(v_j)=16, 1] representing the vector output `v_j` of capsule j in the layer l+1:

def routing(input, b_IJ):
''' The routing algorithm.
Args:
input: A Tensor with [batch_size, num_caps_l=1152, 1, length(u_i)=8, 1]
shape, num_caps_l meaning the number of capsule in the layer l.
Returns:
A Tensor of shape [batch_size, num_caps_l_plus_1, length(v_j)=16, 1]
representing the vector output `v_j` in the layer l+1
Notes:
u_i represents the vector output of capsule i in the layer l, and
v_j the vector output of capsule j in the layer l+1.
'''
  # W: [num_caps_j, num_caps_i, len_u_i, len_v_j]
  W = tf.get_variable('Weight', shape=(1, 1152, 10, 8, 16), dtype=tf.float32,
  initializer=tf.random_normal_initializer(stddev=cfg.stddev))
  # Eq.2, calc u_hat
  # do tiling for input and W before matmul
  # input => [batch_size, 1152, 10, 8, 1]
  # W => [batch_size, 1152, 10, 8, 16]
  input = tf.tile(input, [1, 1, 10, 1, 1])
  W = tf.tile(W, [cfg.batch_size, 1, 1, 1, 1])
  assert input.get_shape() == [cfg.batch_size, 1152, 10, 8, 1]
  
  # in last 2 dims:
  # [8, 16].T x [8, 1] => [16, 1] => [batch_size, 1152, 10, 16, 1]
  u_hat = tf.matmul(W, input, transpose_a=True)
  assert u_hat.get_shape() == [cfg.batch_size, 1152, 10, 16, 1]
  
  # line 3,for r iterations do
  for r_iter in range(cfg.iter_routing):
    with tf.variable_scope('iter_' + str(r_iter)):
      # line 4:
      # => [1, 1152, 10, 1, 1]
      c_IJ = tf.nn.softmax(b_IJ, dim=2)
      c_IJ = tf.tile(c_IJ, [cfg.batch_size, 1, 1, 1, 1])
      assert c_IJ.get_shape() == [cfg.batch_size, 1152, 10, 1, 1]
      # line 5:
      # weighting u_hat with c_IJ, element-wise in the last two dims
      # => [batch_size, 1152, 10, 16, 1]
      s_J = tf.multiply(c_IJ, u_hat)
      # then sum in the second dim, resulting in [batch_size, 1, 10, 16, 1]
      s_J = tf.reduce_sum(s_J, axis=1, keep_dims=True)
      assert s_J.get_shape() == [cfg.batch_size, 1, 10, 16, 16
      # line 6:
      # squash using Eq.1,
      v_J = squash(s_J)
      assert v_J.get_shape() == [cfg.batch_size, 1, 10, 16, 1]
      # line 7:
      # reshape & tile v_j from [batch_size ,1, 10, 16, 1] to [batch_size, 10, 1152, 16, 1]
      # then matmul in the last tow dim: [16, 1].T x [16, 1] => [1, 1], reduce mean in the
      # batch_size dim, resulting in [1, 1152, 10, 1, 1]
      v_J_tiled = tf.tile(v_J, [1, 1152, 1, 1, 1])
      u_produce_v = tf.matmul(u_hat, v_J_tiled, transpose_a=True)
      assert u_produce_v.get_shape() == [cfg.batch_size, 1152, 10, 1, 1]
      b_IJ += tf.reduce_sum(u_produce_v, axis=0, keep_dims=True)
  return(v_J)

6. Now let us review the nonlinear activation squashing function. The input is a 4D vector with the shape [batch_size, num_caps, vec_len, 1] and the output is a 4-D tensor with the same shape as a vector but squashed in the third and fourth dimensions. Given a vector input, the goal is to compute the value represented in the Equation 1 which is shown as follows:

def squash(vector):
'''Squashing function corresponding to Eq. 1
Args:
vector: A 5-D tensor with shape [batch_size, 1, num_caps, vec_len, 1],
Returns:
A 5-D tensor with the same shape as vector but squashed in 4rd and 5th dimensions.
'''
  vec_squared_norm = tf.reduce_sum(tf.square(vector), -2, keep_dims=True)
  scalar_factor = vec_squared_norm / (1 + vec_squared_norm) / tf.sqrt(vec_squared_norm + epsilon)
  vec_squashed = scalar_factor * vector # element-wise
return(vec_squashed)

7. During the previous steps, we have defined what a capsule is, the dynamic routing algorithm among capsules, and the nonlinear squashing function. Now we can define the proper CapsNet. A loss function is built for training and the Adam Optimizer is chosen. The method build_arch(...) defines the CapsNet represented in the following figure:

Note that the paper describes a technique of reconstruction as a regularization method. From the paper:

Then we use this activity vector to reconstruct.

The output of the digit capsule is fed into a decoder consisting of three fully connected layers that model the pixel intensities as described in Fig. 2. We minimize the sum of squared differences between the output of the logistic units and the pixel intensities. We scale down this reconstruction loss by 0.0005 so that it does not dominate the margin loss during training. The method build_arch(..) implemented as follows is also used for creating the Decoder:

#capsNet.py
#
import tensorflow as tf
from config import cfg
from utils import get_batch_data
from capsLayer import CapsLayer
epsilon = 1e-9

class CapsNet(object):
  def __init__(self, is_training=True):
    self.graph = tf.Graph()
    with self.graph.as_default():
      if is_training:
        self.X, self.labels = get_batch_data()
        self.Y = tf.one_hot(self.labels, depth=10, axis=1, dtype=tf.float32)
        self.build_arch()
        self.loss()
        self._summary()

        # t_vars = tf.trainable_variables()
        self.global_step = tf.Variable(0, name='global_step', trainable=False)
        self.optimizer = tf.train.AdamOptimizer()
        self.train_op =    self.optimizer.minimize(self.total_loss, global_step=self.global_step) # var_list=t_vars)

      elif cfg.mask_with_y:
        self.X = tf.placeholder(tf.float32,
          shape=(cfg.batch_size, 28, 28, 1))
        self.Y = tf.placeholder(tf.float32, shape=(cfg.batch_size, 10, 1))
        self.build_arch()
      else:
        self.X = tf.placeholder(tf.float32,
        shape=(cfg.batch_size, 28, 28, 1))
        self.build_arch()
      tf.logging.info('Setting up the main structure')

def build_arch(self):
  
  with tf.variable_scope('Conv1_layer'):
    # Conv1, [batch_size, 20, 20, 256]
    conv1 = tf.contrib.layers.conv2d(self.X, num_outputs=256,
       kernel_size=9, stride=1, 
       padding='VALID')
    assert conv1.get_shape() == [cfg.batch_size, 20, 20, 256]# Primary Capsules layer, return [batch_size, 1152, 8, 1]

  with tf.variable_scope('PrimaryCaps_layer'):
    primaryCaps = CapsLayer(num_outputs=32, vec_len=8,   with_routing=False, layer_type='CONV')
    caps1 = primaryCaps(conv1, kernel_size=9, stride=2)
    assert caps1.get_shape() == [cfg.batch_size, 1152, 8, 1]

  # DigitCaps layer, return [batch_size, 10, 16, 1]
  with tf.variable_scope('DigitCaps_layer'):
    digitCaps = CapsLayer(num_outputs=10, vec_len=16,   with_routing=True, layer_type='FC')
    self.caps2 = digitCaps(caps1)

  # Decoder structure in Fig. 2
  # 1. Do masking, how:
  with tf.variable_scope('Masking'):
    # a). calc ||v_c||, then do softmax(||v_c||)
    # [batch_size, 10, 16, 1] => [batch_size, 10, 1, 1]
    self.v_length = tf.sqrt(tf.reduce_sum(tf.square(self.caps2),
axis=2, keep_dims=True) + epsilon)
    self.softmax_v = tf.nn.softmax(self.v_length, dim=1)
    assert self.softmax_v.get_shape() == [cfg.batch_size, 10, 1, 1]
    # b). pick out the index of max softmax val of the 10 caps
    # [batch_size, 10, 1, 1] => [batch_size] (index)
    self.argmax_idx = tf.to_int32(tf.argmax(self.softmax_v, axis=1))
    assert self.argmax_idx.get_shape() == [cfg.batch_size, 1, 1]
    self.argmax_idx = tf.reshape(self.argmax_idx, shape=(cfg.batch_size, )) .  
    # Method 1.
   if not cfg.mask_with_y:
     # c). indexing
     # It's not easy to understand the indexing process with  argmax_idx
     # as we are 3-dim animal
     masked_v = []
     for batch_size in range(cfg.batch_size):
       v = self.caps2[batch_size][self.argmax_idx[batch_size], :]
       masked_v.append(tf.reshape(v, shape=(1, 1, 16, 1)))
       self.masked_v = tf.concat(masked_v, axis=0)
       assert self.masked_v.get_shape() == [cfg.batch_size, 1, 16, 1]

   # Method 2. masking with true label, default mode
   else:
     self.masked_v = tf.matmul(tf.squeeze(self.caps2), tf.reshape(self.Y, (-1, 10, 1)), transpose_a=True)
     self.v_length = tf.sqrt(tf.reduce_sum(tf.square(self.caps2), axis=2, keep_dims=True) + epsilon)

  # 2. Reconstruct the MNIST images with 3 FC layers
  # [batch_size, 1, 16, 1] => [batch_size, 16] => [batch_size, 512] 
  with tf.variable_scope('Decoder'):
    vector_j = tf.reshape(self.masked_v, shape=(cfg.batch_size, -1))
    fc1 = tf.contrib.layers.fully_connected(vector_j, num_outputs=512)
    assert fc1.get_shape() == [cfg.batch_size, 512]
    fc2 = tf.contrib.layers.fully_connected(fc1, num_outputs=1024)
    assert fc2.get_shape() == [cfg.batch_size, 1024]
    self.decoded = tf.contrib.layers.fully_connected(fc2, num_outputs=784, activation_fn=tf.sigmoid)

8. Another important part defined in the paper is the margin loss function. This is explained in the snippet quote from the paper below (Equation 4) and implemented in the loss(..) method which consists of three losses, the margin loss, the reconstruction loss and the total loss:

def loss(self):
  # 1. The margin loss
  # [batch_size, 10, 1, 1]
  # max_l = max(0, m_plus-||v_c||)^2
  max_l = tf.square(tf.maximum(0., cfg.m_plus - self.v_length))
  # max_r = max(0, ||v_c||-m_minus)^2
  max_r = tf.square(tf.maximum(0., self.v_length - cfg.m_minus))
  assert max_l.get_shape() == [cfg.batch_size, 10, 1, 1]
  
  # reshape: [batch_size, 10, 1, 1] => [batch_size, 10]
  max_l = tf.reshape(max_l, shape=(cfg.batch_size, -1))
  max_r = tf.reshape(max_r, shape=(cfg.batch_size, -1))
  # calc T_c: [batch_size, 10]
  T_c = self.Y
  # [batch_size, 10], element-wise multiply
  L_c = T_c * max_l + cfg.lambda_val * (1 - T_c) * max_r

  self.margin_loss = tf.reduce_mean(tf.reduce_sum(L_c, axis=1))
  
  # 2. The reconstruction loss
  orgin = tf.reshape(self.X, shape=(cfg.batch_size, -1))
  squared = tf.square(self.decoded - orgin)
  self.reconstruction_err = tf.reduce_mean(squared)

  # 3. Total loss
  # The paper uses sum of squared error as reconstruction   error, but we
  # have used reduce_mean in `# 2 The reconstruction loss` to calculate
  # mean squared error. In order to keep in line with the paper,the
  # regularization scale should be 0.0005*784=0.392
  self.total_loss = self.margin_loss + cfg.regularization_scale * self.reconstruction_err

9. In addition, it might be convenient to define a _summary(...) method to report the losses and the accuracy:

#Summary
def _summary(self):
  train_summary = []
  train_summary.append(tf.summary.scalar('train/margin_loss', self.margin_loss))train_summary.append(tf.summary.scalar('train/reconstruction_loss', self.reconstruction_err))
  train_summary.append(tf.summary.scalar('train/total_loss', self.total_loss))
  recon_img = tf.reshape(self.decoded, shape=(cfg.batch_size, 28, 28, 1))
  train_summary.append(tf.summary.image('reconstruction_img', recon_img))
  correct_prediction = tf.equal(tf.to_int32(self.labels), self.argmax_idx)
  self.batch_accuracy = tf.reduce_sum(tf.cast(correct_prediction, tf.float32))
  self.test_acc = tf.placeholder_with_default(tf.constant(0.), shape=[])
  test_summary = []
  test_summary.append(tf.summary.scalar('test/accuracy', self.test_acc))
  self.train_summary = tf.summary.merge(train_summary)
  self.test_summary = tf.summary.merge(test_summary)