for _ in range(self.num_routing):
c = K.softmax(self.bias)
c_expand = K.expand_dims(K.expand_dims(K.expand_dims(c, 2), 2), 0)
outputs = K.sum(c_expand * inputs_hat, 1, keepdims=True)
outputs = squash(outputs)
self.bias = K.update(self.bias, self.bias + K.sum(inputs_hat * outputs, [0, -2, -1]))
// Handling with no routing scenario. Prior bias will always be zero.
if self.num_routing == 0:
c = K.softmax(self.bias)
c_expand = K.expand_dims(K.expand_dims(K.expand_dims(c, 2), 2), 0)
outputs = squash(K.sum(c_expand * inputs_hat, 1, keepdims=True))
return K.reshape(outputs, [-1, self.num_capsule, self.dim_vector])
After Change
_, self.bias, outputs = tf.while_loop(cond, body, loop_vars)
// Routing algorithm V2. Use for iteration. V2 and V1 both work without much difference on performance
assert self.num_routing > 0, "The num_routing should be > 0."
for i in range(self.num_routing):
c = tf.nn.softmax(self.bias, dim=2) // dim=2 is the num_capsule dimension
// outputs.shape=[None, 1, num_capsule, 1, dim_vector]
outputs = squash(K.sum(c * inputs_hat, 1, keepdims=True))
