Custom training step in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Custom training step

Problem:You are training a simple neural network on the MNIST dataset using TensorFlow. The current training uses the built-in model.fit() method.

Current Metrics:Training accuracy: 98%, Validation accuracy: 96%, Training loss: 0.05, Validation loss: 0.08

Issue:You want to customize the training process to add a custom metric and control the training loop manually, but you are unsure how to implement a custom training step correctly.

Your Task

Implement a custom training step function that computes loss, applies gradients, and updates a custom accuracy metric. Train the model for 3 epochs and achieve validation accuracy close to or better than 96%.

Use TensorFlow 2.x eager execution style.

Do not use model.fit() for training; instead, write a custom training loop.

Keep the model architecture the same as the original simple neural network.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

import tensorflow as tf

# Load MNIST dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Normalize and expand dims
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(64)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(64)

# Define simple model
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, activation='relu'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

val_loss = tf.keras.metrics.Mean(name='val_loss')
val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='val_accuracy')

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)

@tf.function
def val_step(images, labels):
    predictions = model(images, training=False)
    v_loss = loss_object(labels, predictions)

    val_loss(v_loss)
    val_accuracy(labels, predictions)

EPOCHS = 3

for epoch in range(EPOCHS):
    train_loss.reset_states()
    train_accuracy.reset_states()
    val_loss.reset_states()
    val_accuracy.reset_states()

    for images, labels in train_ds:
        train_step(images, labels)

    for val_images, val_labels in test_ds:
        val_step(val_images, val_labels)

    print(f'Epoch {epoch + 1}, '
          f'Loss: {train_loss.result():.4f}, '
          f'Accuracy: {train_accuracy.result() * 100:.2f}%, '
          f'Val Loss: {val_loss.result():.4f}, '
          f'Val Accuracy: {val_accuracy.result() * 100:.2f}%')

Replaced model.fit() with a custom training loop using tf.GradientTape.

Created train_step and val_step functions to compute loss, apply gradients, and update metrics.

Used tf.keras.metrics to track loss and accuracy manually.

Added metric reset at the start of each epoch to ensure correct metric calculation.

Results Interpretation

Before: Training accuracy: 98%, Validation accuracy: 96%, Training loss: 0.05, Validation loss: 0.08

After: Training accuracy: ~98%, Validation accuracy: ~96.5%, Training loss: ~0.05, Validation loss: ~0.07

Writing a custom training step lets you control exactly how the model learns. You can add custom metrics or modify the training process while keeping performance similar to built-in methods.

Bonus Experiment

Modify the custom training step to include L2 weight regularization and observe its effect on validation accuracy.

💡 Hint

Add L2 regularization loss to the total loss inside the train_step function by summing model.losses.

Practice

1.

What is the main purpose of a custom training step in TensorFlow?

easy

A. To visualize model predictions during training

B. To automatically generate a model architecture

C. To load datasets from external sources

D. To define exactly how the model learns from each batch of data

Solution

1. Step 1: Understand the role of a training step

   A training step controls how the model updates its weights using data batches.

2. Step 2: Identify what custom training step changes

   Custom training steps let you write your own logic for learning, including gradient calculation and weight updates.

3. Final Answer:

   To define exactly how the model learns from each batch of data -> Option D

4. Quick Check:

   Custom training step = control learning process [OK]

Hint: Custom training step controls learning per batch [OK]

Common Mistakes:

• Thinking it builds model architecture
• Confusing with data loading
• Assuming it only visualizes outputs

2.

Which of the following is the correct way to start recording operations for gradient calculation in a custom training step?

def train_step(inputs):
    # What goes here?
    ...

easy

A. with tf.Variable() as var:

B. with tf.GradientTape() as tape:

C. with tf.Session() as sess:

D. with tf.data.Dataset() as ds:

Solution

1. Step 1: Recall TensorFlow's gradient recording method

   TensorFlow uses tf.GradientTape() to record operations for automatic differentiation.

2. Step 2: Match correct syntax for gradient recording

   The correct syntax to start recording is with tf.GradientTape() as tape:.

3. Final Answer:

   with tf.GradientTape() as tape: -> Option B

4. Quick Check:

   Gradient recording uses tf.GradientTape [OK]

Hint: Use tf.GradientTape() to record gradients [OK]

Common Mistakes:

• Using tf.Session() which is for old TF versions
• Confusing tf.Variable() with gradient recording
• Using tf.data.Dataset() incorrectly here

3.

What will be the output of the following custom training step code snippet?

import tensorflow as tf

model = tf.keras.Sequential([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = tf.reduce_mean(tf.square(y - predictions))
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

x = tf.constant([[1.0], [2.0]])
y = tf.constant([[3.0], [5.0]])
loss_value = train_step(x, y).numpy()
print(round(loss_value, 2))

medium

A. 1.0

B. 9.0

C. 4.0

D. 0.0

Solution

1. Step 1: Calculate initial predictions and loss

   Model starts with random weights near zero, so predictions ~0. Loss = mean squared error between y and predictions = mean of [9, 25] = 17.

2. Step 2: After one gradient update with learning rate 0.1, loss reduces

   One step reduces loss roughly to around 4.0 (not zero yet).

3. Final Answer:

   4.0 -> Option C

4. Quick Check:

   One step reduces loss from ~17 to ~4 [OK]

Hint: One step reduces loss but not to zero [OK]

Common Mistakes:

• Assuming loss becomes zero immediately
• Ignoring initial random weights
• Confusing loss with prediction values

4.

Identify the error in this custom training step code and select the fix:

def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x)
        loss = tf.reduce_mean(tf.square(y - predictions))
    gradients = tape.gradient(loss)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

medium

A. Pass model.trainable_variables as second argument to tape.gradient

B. Remove tf.GradientTape() context manager

C. Change loss calculation to tf.reduce_sum

D. Call optimizer.minimize instead of apply_gradients

Solution

1. Step 1: Check tape.gradient usage

   tape.gradient requires two arguments: loss and variables to differentiate with respect to.

2. Step 2: Identify missing argument

   Code misses the second argument (model.trainable_variables), causing error.

3. Final Answer:

   Pass model.trainable_variables as second argument to tape.gradient -> Option A

4. Quick Check:

   tape.gradient(loss, variables) needs variables [OK]

Hint: Always pass variables to tape.gradient [OK]

Common Mistakes:

• Forgetting second argument in tape.gradient
• Removing GradientTape context
• Changing loss function unnecessarily

5.

You want to track both loss and accuracy in a custom training step for a classification model. Which code snippet correctly updates and returns these metrics?

loss_metric = tf.keras.metrics.Mean()
accuracy_metric = tf.keras.metrics.SparseCategoricalAccuracy()

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    # What goes here?
    ...

hard

A. loss_metric.update_state(loss); accuracy_metric.update_state(y, logits); return loss_metric.result(), accuracy_metric.result()

B. loss_metric = loss; accuracy_metric = accuracy(y, logits); return loss_metric, accuracy_metric

C. loss_metric.add(loss); accuracy_metric.add(y, logits); return loss_metric, accuracy_metric

D. loss_metric.reset_states(); accuracy_metric.reset_states(); return loss_metric.result(), accuracy_metric.result()

Solution

1. Step 1: Understand metric update methods

   TensorFlow metrics use update_state() to add new values for tracking.

2. Step 2: Correctly update and return metric results

   Use loss_metric.update_state(loss) and accuracy_metric.update_state(y, logits), then return their result().

3. Final Answer:

   loss_metric.update_state(loss); accuracy_metric.update_state(y, logits); return loss_metric.result(), accuracy_metric.result() -> Option A

4. Quick Check:

   Use update_state() and result() for metrics [OK]

Hint: Use update_state() then result() to track metrics [OK]

Common Mistakes:

• Assigning metric variables instead of updating
• Using non-existent add() method
• Resetting metrics inside training step