Now we will look at how the dynamic routing algorithm is implemented. We use variable names of the same notations that we learned in the dynamic routing algorithm, so that we can easily follow the steps. We will look at each line in our function step by step. You can also check the complete code on GitHub, at http://bit.ly/2HQqDEZ.

First, define the function called dynamic_routing, which takes the previous capsules, ui, coupling coefficients, bij, and number of routing iterations, num_routing as inputs as follows:

def dynamic_routing(ui, bij, num_routing=10):

Initialize the wij weights by drawing from a random normal distribution, and initialize biases with a constant value:

    wij = tf.get_variable('Weight', shape=(1, 1152, 160, 8, 1), dtype=tf.float32,

                        initializer=tf.random_normal_initializer(0.01))

    biases = tf.get_variable('bias', shape=(1, 1, 10, 16, 1))

Define the primary capsules ui (tf.tile replicates the tensor n times):

    ui = tf.tile(ui, [1, 1, 160, 1, 1])

Compute the prediction vector, , as follows:

    u_hat = tf.reduce_sum(wij * ui, axis=3, keep_dims=True)

Reshape the prediction vector:

    u_hat = tf.reshape(u_hat, shape=[-1, 1152, 10, 16, 1])

Stop gradient computation in the prediction vector:

    u_hat_stopped = tf.stop_gradient(u_hat, name='stop_gradient')

Perform dynamic routing for a number of routing iterations, as follows:

    for r in range(num_routing):

        with tf.variable_scope('iter_' + str(r)):
            
            #step 1
            cij = tf.nn.softmax(bij, dim=2)
            
            #step 2
            if r == num_routing - 1:

                sj = tf.multiply(cij, u_hat)

                sj = tf.reduce_sum(sj, axis=1, keep_dims=True) + biases

                vj = squash(sj)

            elif r < num_routing - 1: 

                sj = tf.multiply(cij, u_hat_stopped)

                sj = tf.reduce_sum(sj, axis=1, keep_dims=True) + biases

                vj = squash(sj)

                vj_tiled = tf.tile(vj, [1, 1152, 1, 1, 1])

                coupling_coeff = tf.reduce_sum(u_hat_stopped * vj_tiled, axis=3, keep_dims=True)

                #step 3
                bij += coupling_coeff
   return vj