def M_step(log_R, log_activation, vote, lambda_val=0.01):
R_shape = tf.shape(log_R)
log_R = log_R + log_activation
R_sum_i = cl.reduce_sum(tf.exp(log_R), axis=-3, keepdims=True)
log_normalized_R = log_R - \
tf.reduce_logsumexp(log_R, axis=-3, keepdims=True)
pose = cl.reduce_sum(vote * tf.exp(log_normalized_R),
axis=-3,
keepdims=True)
log_var = tf.reduce_logsumexp(log_normalized_R +
cl.log(tf.square(vote - pose)),
axis=-3,
keepdims=True)
shape = [1 for i in range(len(pose.shape) - 2)] + [pose.shape[-2], 1]
beta_v = tf.Variable(initial_value=tf.keras.initializers.TruncatedNormal(
mean=15., stddev=3.)(shape),
name="beta_v",
shape=shape)
cost = R_sum_i * (beta_v + 0.5 * log_var)
beta_a = tf.Variable(initial_value=tf.keras.initializers.TruncatedNormal(
mean=100.0, stddev=10)(shape),
name="beta_a",
shape=shape)
cost_sum_h = cl.reduce_sum(cost, axis=-1, keepdims=True)
logit = lambda_val * (beta_a - cost_sum_h)
log_activation = tf.math.log_sigmoid(logit)
return (pose, log_var, log_activation)
def _body(i, logits, poses):
if self.leaky:
route = _leaky_routing(logits)
else:
route = tf.nn.softmax(logits, axis=-3)
if self.use_bias:
preactivate = cl.reduce_sum(
route * votes, axis=-4, keepdims=True) + self.biases
else:
preactivate = cl.reduce_sum(route * votes,
axis=-4,
keepdims=True)
pose = cl.ops.squash(preactivate, axis=squash_on)
poses = poses.write(i, pose)
if vote_shape[-1] == 1:
distances = cl.matmul(votes, pose, transpose_a=True)
else:
diff = votes - pose
distances = tf.linalg.trace(cl.matmul(diff,
diff))[..., tf.newaxis,
tf.newaxis]
logits += distances
return (i + 1, logits, poses)
def fcm_loss(inputs, probs, num_iters=5, m=2):
"""
Args:
inputs: [num_samples, 1, x_dims, 1]
probs: [num_samples, num_clusters, 1, 1].
"""
# centers: [1, num_clusters, x_dims]
k = 1
b = 0
weights = []
delta_probs = []
for i in range(num_iters):
probs_m = tf.pow(probs, m)
centers = cl.reduce_sum(probs_m * inputs, axis=0, keepdims=True) / cl.reduce_sum(probs_m, axis=0, keepdims=True)
# distance matrix with shape [num_samples, num_clusters, 1, 1]
distance_matrix = cl.norm(inputs - centers, axis=(2, 3), keepdims=True)
distance_matrix = tf.pow(distance_matrix, 2 / (m - 1))
probs_plus = 1 / (distance_matrix / cl.reduce_sum(distance_matrix, axis=1, keepdims=True))
delta_probs.append(tf.norm(probs_plus - probs))
weights.append(tf.exp(tf.cast(k * i + b, tf.float32)))
probs = probs_plus
weights = tf.stack(weights, axis=0)
delta_probs = tf.stack(delta_probs, axis=0)
loss = tf.reduce_sum(weights * delta_probs) / tf.reduce_sum(weights)
return loss
def dynamicRouting_v1(votes,
num_routing=3,
use_bias=True,
leaky=True):
""" Dynamic routing algorithm.
See [Sabour et al., 2017](https://arxiv.org/abs/1710.09829).
Args:
votes: A 5-D or 7-D tensor with shape [batch_size, ..., in_channels/num_inputs, num_outputs] + out_caps_dims.
num_routing: Integer, number of routing iterations.
use_bias: Boolean, whether the layer uses a bias.
leaky: Boolean, whether the algorithm uses leaky routing.
Returns:
poses: A 4-D or 6-D tensor.
probs: A 2-D or 4-D tensor.
"""
vote_shape = cl.shape(votes)
logit_shape = vote_shape[:-2] + [1, 1]
logits = tf.fill(logit_shape, 0.0)
squash_on = -2 if vote_shape[-1] == 1 else [-2, -1]
if use_bias:
bias_shape = [1 for i in range(len(vote_shape) - 3)] + vote_shape[-3:]
biases = tf.get_variable("biases",
bias_shape,
initializer=tf.constant_initializer(0.1),
dtype=tf.float32)
vote_stopped = tf.stop_gradient(votes, name="stop_gradient")
for i in range(num_routing):
with tf.variable_scope("iter_" + str(i)):
if leaky:
route = _leaky_routing(logits)
else:
route = cl.softmax(logits, axis=-3)
if i == num_routing - 1:
if use_bias:
preactivate = cl.reduce_sum(tf.multiply(route, votes), axis=-4, keepdims=True) + biases
else:
preactivate = cl.reduce_sum(tf.multiply(route, votes), axis=-4, keepdims=True)
poses = cl.ops.squash(preactivate, axis=squash_on)
else:
if use_bias:
preactivate = cl.reduce_sum(tf.multiply(route, vote_stopped), axis=1, keepdims=True) + biases
else:
preactivate = cl.reduce_sum(tf.multiply(route, vote_stopped), axis=1, keepdims=True)
poses = cl.ops.squash(preactivate, axis=squash_on)
logits += cl.reduce_sum(vote_stopped * poses, axis=-4, keepdims=True)
poses = tf.squeeze(poses, axis=-4)
probs = tf.norm(poses, axis=(-2, -1))
return(poses, probs)
def E_step(pose, log_var, log_activation, vote):
normalized_vote = cl.divide(tf.square(vote - pose), 2 * tf.exp(log_var))
log_probs = normalized_vote + cl.log(2 * np.pi) + log_var
log_probs = -0.5 * cl.reduce_sum(log_probs, axis=-1, keepdims=True)
log_activation_logit = log_activation + log_probs
log_activation_logit = log_probs
log_R = log_activation_logit - tf.reduce_logsumexp(log_activation_logit, axis=-2, keepdims=True)
return log_R
def M_step(self, log_R, log_activation, vote, lambda_val=0.01):
log_R = log_R + log_activation
R_sum_i = cl.reduce_sum(tf.exp(log_R), axis=-3, keepdims=True)
log_normalized_R = log_R - \
tf.reduce_logsumexp(log_R, axis=-3, keepdims=True)
pose = cl.reduce_sum(vote * tf.exp(log_normalized_R),
axis=-3,
keepdims=True)
log_var = tf.reduce_logsumexp(log_normalized_R +
cl.log(tf.square(vote - pose)),
axis=-3,
keepdims=True)
cost = R_sum_i * (self.beta_v + 0.5 * log_var)
cost_sum_h = cl.reduce_sum(cost, axis=-1, keepdims=True)
logit = lambda_val * (self.beta_a - cost_sum_h)
log_activation = tf.math.log_sigmoid(logit)
return (pose, log_var, log_activation)
def M_step(log_R, log_activation, vote, lambda_val=0.01):
R_shape = tf.shape(log_R)
log_R = log_R + log_activation
R_sum_i = cl.reduce_sum(tf.exp(log_R), axis=-3, keepdims=True)
log_normalized_R = log_R - tf.reduce_logsumexp(log_R, axis=-3, keepdims=True)
pose = cl.reduce_sum(vote * tf.exp(log_normalized_R), axis=-3, keepdims=True)
log_var = tf.reduce_logsumexp(log_normalized_R + cl.log(tf.square(vote - pose)), axis=-3, keepdims=True)
beta_v = tf.get_variable('beta_v',
shape=[1 for i in range(len(pose.shape) - 2)] + [pose.shape[-2], 1],
initializer=tf.truncated_normal_initializer(mean=15., stddev=3.))
cost = R_sum_i * (beta_v + 0.5 * log_var)
beta_a = tf.get_variable('beta_a',
shape=[1 for i in range(len(pose.shape) - 2)] + [pose.shape[-2], 1],
initializer=tf.truncated_normal_initializer(mean=100.0, stddev=10))
cost_sum_h = cl.reduce_sum(cost, axis=-1, keepdims=True)
logit = lambda_val * (beta_a - cost_sum_h)
log_activation = tf.log_sigmoid(logit)
return(pose, log_var, log_activation)
def spread_loss(labels, logits, margin, regularizer=None):
"""
Args:
labels: [batch_size, num_label].
logits: [batch_size, num_label].
margin: Integer or 1-D Tensor.
regularizer: use regularization.
Returns:
loss: Spread loss.
"""
a_target = cl.reduce_sum(labels * logits, axis=1, keepdims=True)
dist = (1 - labels) * margin - (a_target - logits)
dist = tf.pow(tf.maximum(0., dist), 2)
loss = tf.reduce_mean(tf.reduce_sum(dist, axis=-1))
if regularizer is not None:
loss += tf.reduce_mean(regularizer)
return (loss)
