def eval_step(model, images, labels, metrics_used):
# Note: wrap whole test step, otherwise you end up with a tape on the CPU.
with tf.device('/gpu:*'):
features, predictions = model(images, training=False)
metrics.update_state(metrics_used, labels, predictions)
return
def inner_loop(images_aug_1, images_aug_2, images, labels):
# copy weights for each image
# online_model.set_weights(target_model.get_weights()) # slow
for k in range(0, len(online_model.weights)):
if not online_model.weights[k].dtype == tf.bool:
online_model.weights[k].assign(target_model.weights[k])
# acc without inner update
_, _, predictions = online_model(
images[:1, :, :, :],
training=False) # only one image since repetition
metrics.update_state(metrics_test[0], labels[:1, :], predictions)
# inner update and acc
for k in range(num_test_steps):
# calc target
# Get targets
_, tar1, _ = target_model(images_aug_1,
use_predictor=False,
training=True)
_, tar2, _ = target_model(images_aug_2,
use_predictor=False,
training=True)
# Perform inner optimization
with tf.GradientTape(persistent=False) as test_tape:
_, prediction1, _ = online_model(images_aug_1,
use_predictor=True,
training=True)
_, prediction2, _ = online_model(images_aug_2,
use_predictor=True,
training=True)
# Calc byol loss
loss1 = byol_loss_fn(prediction1, tf.stop_gradient(tar2))
loss2 = byol_loss_fn(prediction2, tf.stop_gradient(tar1))
loss = tf.reduce_mean(loss1 + loss2)
gradients = test_tape.gradient(loss,
online_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, online_model.trainable_variables))
# get predictions for test acc
_, _, predictions = online_model(
images[:1, :, :, :],
training=False) # only one image since repetition
metrics.update_state(metrics_test[k + 1], labels[:1, :],
predictions)
return 0
def train_step(model,
images,
labels,
optimizer,
loss_obj,
metric_loss,
metrics_used,
epoch_tf,
b_verbose=False):
logging.info(
f'Trace indicator - train epoch - eager mode: {tf.executing_eagerly()}.'
)
# Note: wrap whole training step, otherwise you end up with a tape on the CPU.
with tf.device('/gpu:*'):
with tf.GradientTape() as tape:
features, predictions = model(
images, training=True) # graph is build on first model.call()
loss = loss_obj(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
# update metrics
metric_loss.update_state(loss)
metrics.update_state(metrics_used, labels, predictions)
if b_verbose:
tf.print("Training loss for epoch:",
epoch_tf + 1,
" and step: ",
optimizer.iterations,
" - ",
loss,
output_stream=sys.stdout)
return
def train_step(online_model,
target_model,
image_aug1,
image_aug2,
image_sup,
labels,
optimizer,
ce_loss_obj,
metric_loss,
metric_ce_loss,
metric_byol_loss,
metrics_used,
epoch_tf,
beta_tf,
gamma_byol,
b_verbose=False):
logging.info(
f'Trace indicator - train epoch - eager mode: {tf.executing_eagerly()}.'
)
# Note: wrap whole training step, otherwise you end up with a tape on the CPU.
with tf.device('/gpu:*'):
# Get targets (not backproped for optimization)
_, target_aug1, _ = target_model(image_aug1,
use_predictor=False,
training=True)
_, target_aug2, _ = target_model(image_aug2,
use_predictor=False,
training=True)
with tf.GradientTape() as tape:
_, online_aug1, _ = online_model(image_aug1,
use_predictor=True,
training=True)
_, online_aug2, _ = online_model(image_aug2,
use_predictor=True,
training=True)
_, _, predictions = online_model(image_sup,
use_predictor=True,
training=True)
# Calculate losses
byol_loss = tf.reduce_mean(
byol_loss_fn(online_aug1, tf.stop_gradient(target_aug2)) +
byol_loss_fn(online_aug2, tf.stop_gradient(target_aug1)))
ce_loss = ce_loss_obj(labels, predictions)
loss = ce_loss + gamma_byol * byol_loss
gradients = tape.gradient(loss, online_model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, online_model.trainable_variables))
# Update target model weights with MAVG
for k in range(0, len(target_model.weights)):
if not target_model.weights[k].dtype == tf.bool:
target_model.weights[k].assign(target_model.weights[k] *
beta_tf +
online_model.weights[k] *
(1 - beta_tf))
# update metrics
metric_loss.update_state(loss)
metric_ce_loss.update_state(ce_loss)
metric_byol_loss.update_state(byol_loss)
metrics.update_state(metrics_used, labels, predictions)
if b_verbose:
tf.print("Training loss for epoch:",
epoch_tf + 1,
" and step: ",
optimizer.iterations,
" - loss: ",
loss,
" - ce_loss: ",
ce_loss,
" - byol_loss: ",
byol_loss,
output_stream=sys.stdout)
return
def eval_step(self, images, labels):
_, _, predictions = self.target_model(images)
metrics.update_state(self.metrics_eval, labels, predictions)
return 0