def update_weights(optimizer: tf.keras.optimizers.Optimizer,
network: BaseNetwork, batch, weight_decay: float):
loss = 0
for image, actions, targets in batch:
# initial step, from the real observation
value, reward, policy_logits, hidden_state = network.initial_inference(
image)
predictions = [(1., value, reward, policy_logits)]
# recurrent steps, from action and previous hidden state
for action in actions:
value, reward, policy_logits, hidden_state = network.recurrent_inference(
hidden_state, action)
predictions.append((1., len(actions), value, reward, policy_logits))
hidden_state = scale_gradient(hidden_state, 0.5)
for prediction, target in zip(predictions, targets):
gradient_scale, value, reward, policy_logits = prediction
target_value, target_reward, target_policy = target
l = (scalar_loss(value, target_value) +
scalar_loss(reward, target_reward) +
tf.nn.softmax_cross_entropy_with_logits(logits=policy_logits,
labels=target_policy))
loss += l
for weights in network.get_weights():
loss += weight_decay * tf.nn.l2_loss(weights)
optimizer.minimize(loss=loss, var_list=network.cb_get_variables())
def build(inputs: tf.Tensor,
labels: tf.Tensor,
classes,
optimizer: tf.keras.optimizers.Optimizer = None,
global_step=None,
dtype: tf.DType = tf.float32):
"""
Build model
:param inputs: Input Tensor. For MNIST with shape [batch, 28, 28, 1]
:param labels: Label Tensor with shape [batch]
:param classes: Number of classes
:param optimizer: An optimizer of choice. Default is Adam.
:param global_step: Optionally a global step Variable
:param dtype: Desired data type
:return:
"""
training = tf.placeholder_with_default(False, shape=(), name='training')
if optimizer is None:
optimizer = tf.train.AdamOptimizer()
# Encoder
encoder_out, probabilities = encoder(inputs, classes)
predictions = tf.argmax(probabilities,
axis=-1,
output_type=labels.dtype,
name="predictions")
eq = tf.cast(tf.equal(predictions, labels), dtype=dtype)
correct_predictions = tf.reduce_sum(eq, name="correct_predictions")
false_predictions = tf.reduce_sum(1 - eq, name="false_predictions")
accuracy = tf.reduce_mean(eq, name="accuracy")
# Decoder
encoder_out_masked_flat = tc.layers.label_mask(encoder_out, labels,
predictions, training)
decoder_out = decoder(
encoder_out_masked_flat,
inputs.shape.as_list()[1:]) # change this line for dynamic shapes
# Loss
loss = tc.losses.capsule_net_loss(inputs,
encoder_out,
decoder_out,
labels,
m_minus=.1,
m_plus=.9,
alpha=.0005)
train_op = optimizer.minimize(loss, global_step)
return train_op, training, correct_predictions, false_predictions, accuracy, loss