Function train_step

examples/basic_tutorials/tutorial_mnist_siamese.py:104–113

(X_batch, y_batch)

@tf.function
def train_step(X_batch, y_batch):
    with tf.GradientTape() as tape:
        _out1, _out2 = model([X_batch[:, 0, :], X_batch[:, 1, :]])
        _loss = contrastive_loss(y_batch, _out1, _out2)

    grad = tape.gradient(_loss, train_weights)
    optimizer.apply_gradients(zip(grad, train_weights))

    _acc = compute_accuracy(y_batch, _out1, _out2)
    return _loss, _acc


# begin training