def capsule_discriminator(x_image, reuse=True):
""" capsule network as discriminator
"""
if reuse:
tf.get_variable_scope().reuse_variables()
x = tf.reshape(x_image, [-1, 28, 28, 1])
conv1 = tf.layers.conv2d(inputs=x,
filters=256,
kernel_size=[9, 9],
padding="valid",
activation=tf.nn.relu,
name="d_ReLU_Conv1")
conv1 = tf.expand_dims(conv1, axis=-2)
# Convolutional capsules
with tf.variable_scope('d_PrimaryCaps'):
primary_caps = capsule.conv2d(conv1, 32, 8, [9, 9], strides=(2, 2))
primary_caps = tf.reshape(primary_caps, [
-1, primary_caps.shape[1].value * primary_caps.shape[2].value * 32,
8
])
# Fully Connected capsules with routing by agreement. Binary classifier.
with tf.variable_scope('d_DigitCaps'):
digit_caps = capsule.dense(primary_caps,
1,
16,
iter_routing=3,
learn_coupling=False,
mapfn_parallel_iterations=16)
# The length of the capsule activation vectors.
length = tf.sqrt(tf.reduce_sum(tf.square(digit_caps), axis=1),
name="Length")
return length
def caps_model_fn(features, labels, mode):
hooks = []
train_log_dict = {}
"""Model function for CNN."""
# Input Layer
# Reshape X to 4-D tensor: [batch_size, width, height, channels]
# Fashion MNIST images are 28x28 pixels, and have one color channel
input_layer = tf.reshape(features["x"], [-1, 28, 28, 1])
# A little bit cheaper version of the capsule network in: Dynamic Routing Between Capsules
# Std. convolutional layer
conv1 = tf.layers.conv2d(inputs=input_layer,
filters=256,
kernel_size=[9, 9],
padding="valid",
activation=tf.nn.relu,
name="ReLU_Conv1")
conv1 = tf.expand_dims(conv1, axis=-2)
# Convolutional capsules, no routing as the dimension of the units of previous layer is one
primarycaps = caps.conv2d(conv1,
32,
8, [9, 9],
strides=(2, 2),
name="PrimaryCaps")
primarycaps = tf.reshape(
primarycaps,
[-1, primarycaps.shape[1].value * primarycaps.shape[2].value * 32, 8])
# Fully connected capsules with routing by agreement
digitcaps = caps.dense(primarycaps,
10,
16,
iter_routing=iter_routing,
learn_coupling=learn_coupling,
mapfn_parallel_iterations=mapfn_parallel_iterations,
name="DigitCaps")
# The length of the capsule activation vectors encodes the probability of an entity being present
lengths = tf.sqrt(tf.reduce_sum(tf.square(digitcaps), axis=2) + epsilon,
name="Lengths")
# Predictions for (PREDICTION mode)
predictions = {
# Generate predictions (for PREDICT and EVAL mode)
"classes": tf.argmax(lengths, axis=1),
"probabilities": tf.nn.softmax(lengths, name="Softmax")
}
if regularization:
masked_digitcaps_pred = mask_one(digitcaps,
lengths,
is_predicting=True)
with tf.variable_scope(tf.get_variable_scope()):
reconstruction_pred = decoder_nn(masked_digitcaps_pred)
predictions["reconstruction"] = reconstruction_pred
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
# Calculate Loss (for both TRAIN and EVAL modes)
onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)
m_loss = margin_loss(onehot_labels, lengths)
train_log_dict["margin loss"] = m_loss
tf.summary.scalar("margin_loss", m_loss)
if regularization:
masked_digitcaps = mask_one(digitcaps, onehot_labels)
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
reconstruction = decoder_nn(masked_digitcaps)
rec_loss = reconstruction_loss(input_layer, reconstruction)
train_log_dict["reconstruction loss"] = rec_loss
tf.summary.scalar("reconstruction_loss", rec_loss)
loss = m_loss + lambda_reg * rec_loss
else:
loss = m_loss
# Configure the Training Op (for TRAIN mode)
if mode == tf.estimator.ModeKeys.TRAIN:
# Logging hook
train_log_dict["accuracy"] = tf.metrics.accuracy(
labels=labels, predictions=predictions["classes"])[1]
logging_hook = tf.train.LoggingTensorHook(
train_log_dict, every_n_iter=config.save_summary_steps)
# Summary hook
summary_hook = tf.train.SummarySaverHook(
save_steps=config.save_summary_steps,
output_dir=model_dir,
summary_op=tf.summary.merge_all())
hooks += [logging_hook, summary_hook]
global_step = tf.train.get_or_create_global_step()
learning_rate = tf.train.exponential_decay(start_lr, global_step,
decay_steps, decay_rate)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss=loss,
global_step=tf.train.get_global_step())
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op,
training_hooks=hooks)
# Add evaluation metrics (for EVAL mode)
eval_metric_ops = {
"accuracy":
tf.metrics.accuracy(labels=labels, predictions=predictions["classes"])
}
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
eval_metric_ops=eval_metric_ops)