def test_shape_2(self):
with self.test_session():
s = helpers.broadcast(
tf.convert_to_tensor(np.random.rand(10, 1, 15, 1, 15)),
tf.convert_to_tensor(np.random.rand(1, 10, 1, 15, 3)),
axis=[0, 1],
broadcast_b=False)
a, b = s[0].eval(), s[1].eval()
self.assertAllCloseAccordingToType(a.shape, [10, 10, 15, 1, 15])
self.assertAllCloseAccordingToType(b.shape, [1, 10, 1, 15, 3])
def _routing(predictions, coupling_logits, routing_iterations):
"""
Performs capsule routing.
Parameters
----------
predictions: tf.Tensor
Lower level capsule predictions `(batch, num_caps[l], num_caps[l-1], 1, units[l])`
coupling_logits: tf.Tensor
Starting coupling logits `b` `(batch, num_caps[l], num_caps[l-1], 1)`
routing_iterations: int
Number of routing iterations
Returns
-------
tf.Tensor
Resulting capsule values. `(batch, num_caps[l], units[l])`
"""
with tf.name_scope("routing"):
for i in range(routing_iterations):
coupling_coeffs = _coupling_coefficients(coupling_logits)
out = tf.reduce_sum(tf.multiply(
predictions, tf.expand_dims(coupling_coeffs, -1)),
axis=-3)
out = helpers.squash(out) # (batch, num_caps[l], 1, units[l])
if i < routing_iterations - 1:
# expand to (batch, num_caps[l], 1, units[l], 1)
out = tf.expand_dims(out, -1)
# broadcast to (batch, num_caps[l], num_caps[l-1], units[l], 1)
out, _ = helpers.broadcast(out,
predictions,
axis=-3,
broadcast_b=False)
# get logit update (batch, num_caps[l], num_caps[l-1], 1, 1)
logits_update = tf.matmul(predictions, out)
logits_update = tf.squeeze(
logits_update,
[-1]) # (batch, num_caps[l], num_caps[l-1], 1)
# update coupling_logits
coupling_logits = tf.add(coupling_logits, logits_update)
return tf.squeeze(
out, axis=-2) # squeeze to get (batch, num_caps[l], units[l])
def _prediction_vectors(inputs, weights):
"""
Computes prediction vectors u_hat.
Inputs should be of shape `(batch, num_caps[l-1], units[l-1])`.
Weights should be of shape `(num_caps[l], num_caps[l-1], units[l-1], units[l])`
Per each capsule[l] computes prediction vector for each lower level capsule.
Outputs tensor of predictions `(batch, num_caps[l], num_caps[l-1], 1, units[l])`.
Parameters
----------
inputs: tf.Tensor
Lower level capsules tensor.
weights: tf.Tensor
Weights used to compute predictions.
Returns
-------
tf.Tensor
Predictions u_hat.
"""
with tf.name_scope("prediction_vectors"):
# take inputs (batch,num_caps[l-1],units[l-1])
# transform into (batch, 1, num_caps[l-1], 1, units[l-1])
inputs = tf.expand_dims(inputs, -2) # add dim before units
inputs = tf.expand_dims(inputs, 1) # add dim after batch
# take weights (num_caps[l], num_caps[l-1], units[l-1], units[l])
# transform into (1, num_caps[l], num_caps[l-1], units[l-1], units[l])
weights = tf.expand_dims(weights, 0)
# inputs (batch, num_caps[l], num_caps[l-1], 1, units[l-1])
# weights (batch, num_caps[l], num_caps[l-1], units[l-1], units[l])
inputs, weights = helpers.broadcast(inputs, weights, axis=[0, 1])
return tf.matmul(inputs, weights)
def test_shape_legacy(self):
with self.test_session():
s = helpers.broadcast(np.random.rand(10, 1, 15), np.random.rand(1, 15, 3), use_legacy=True)
a, b = s[0].eval(), s[1].eval()
self.assertAllCloseAccordingToType(a.shape, [10, 1, 15])
self.assertAllCloseAccordingToType(b.shape, [10, 15, 3])