Why custom loops offer full control in TensorFlow - Experiment to Prove It

Experiment - Why custom loops offer full control

Problem:You want to train a neural network on the MNIST dataset but need more control over the training process than what built-in methods offer.

Current Metrics:Training accuracy: 98%, Validation accuracy: 96%, Loss: 0.07

Issue:Using model.fit() hides the training steps, limiting customization like custom loss calculations, gradient clipping, or complex training logic.

Your Task

Rewrite the training process using a custom training loop to gain full control over each training step and verify that the model still trains correctly with similar accuracy.

Use TensorFlow 2.x eager execution.

Do not use model.fit() or model.evaluate() methods.

Keep the model architecture the same as the original.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models, datasets

# Load and preprocess data
(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()
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 model
model = models.Sequential([
    layers.Conv2D(32, 3, activation='relu'),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(10)
])

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

train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
val_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()

# Custom training loop
for epoch in range(5):
    print(f'Start of epoch {epoch + 1}')
    # Training loop
    for step, (x_batch_train, y_batch_train) in enumerate(train_ds):
        with tf.GradientTape() as tape:
            logits = model(x_batch_train, training=True)
            loss_value = loss_fn(y_batch_train, logits)
        grads = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))
        train_acc_metric.update_state(y_batch_train, logits)

        if step % 200 == 0:
            print(f'Training loss at step {step}: {loss_value:.4f}')

    train_acc = train_acc_metric.result()
    print(f'Training accuracy over epoch: {train_acc:.4f}')
    train_acc_metric.reset_state()

    # Validation loop
    for x_batch_val, y_batch_val in test_ds:
        val_logits = model(x_batch_val, training=False)
        val_acc_metric.update_state(y_batch_val, val_logits)
    val_acc = val_acc_metric.result()
    val_acc_metric.reset_state()
    print(f'Validation accuracy: {val_acc:.4f}')

Replaced model.fit() with a manual loop over batches.

Used tf.GradientTape() to compute gradients explicitly.

Manually applied gradients using optimizer.apply_gradients().

Calculated and printed loss and accuracy inside the loop.

Added separate validation loop to measure accuracy.

Results Interpretation

Before: Used model.fit() with training accuracy 98% and validation accuracy 96%. Training steps were hidden.

After: Custom loop achieved similar accuracy and loss but allowed printing loss every 200 steps and manual control over training.

Custom training loops give you full control over every step of training. This lets you add custom logic, debug easily, and experiment with advanced techniques that built-in methods hide.

Bonus Experiment

Add gradient clipping to the custom training loop to prevent exploding gradients and observe its effect on training stability.

💡 Hint

Use tf.clip_by_norm() on gradients before applying them with optimizer.

Practice

1. Why do custom training loops in TensorFlow give you full control over the training process?

easy

A. Because they require no coding and run faster by default

B. Because TensorFlow automatically handles all training details for you

C. Because custom loops use less memory than built-in training methods

D. Because you manually define each step like forward pass, loss calculation, and backpropagation

Solution

1. Step 1: Understand what custom loops do

   Custom loops let you write code for each training step, like calculating predictions and losses.

2. Step 2: Compare with automatic training

   Built-in methods hide these steps, but custom loops expose them for full control.

3. Final Answer:

   Because you manually define each step like forward pass, loss calculation, and backpropagation -> Option D

4. Quick Check:

   Manual step control = Full control [OK]

Hint: Custom loops mean you write every training step yourself [OK]

Common Mistakes:

• Thinking TensorFlow handles all details automatically
• Assuming custom loops run faster by default
• Believing custom loops require no coding

2. Which of the following is the correct syntax to start a custom training loop in TensorFlow?

easy

A. train_loop(num_epochs):

B. while epoch < num_epochs {}

C. for epoch in range(num_epochs):

D. loop(epoch in num_epochs)

Solution

1. Step 1: Identify Python loop syntax

   In Python, a for loop over a range is written as for variable in range(...):.

2. Step 2: Match with TensorFlow custom loop

   Custom loops use standard Python loops like for epoch in range(num_epochs): to iterate epochs.

3. Final Answer:

   for epoch in range(num_epochs): -> Option C

4. Quick Check:

   Python for loop syntax = for epoch in range(num_epochs): [OK]

Hint: Custom loops use normal Python for loops for epochs [OK]

Common Mistakes:

• Using invalid loop syntax like curly braces
• Confusing function calls with loops
• Using incorrect keywords like 'loop' or 'while' without colon

3. Consider this snippet of a TensorFlow custom training loop:

for epoch in range(2):
    for x, y in dataset:
        with tf.GradientTape() as tape:
            predictions = model(x)
            loss = loss_fn(y, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
print(loss.numpy())

What will be printed after training?

medium

A. The loss value from the last batch of the second epoch

B. The total loss summed over all batches

C. The initial loss before training

D. The average loss over all batches

Solution

1. Step 1: Understand the loop structure

   The loop runs 2 epochs, updating loss each batch. The printed loss is from the last batch processed.

2. Step 2: Identify what loss.numpy() prints

   It prints the loss tensor's value from the last batch, not a sum or average.

3. Final Answer:

   The loss value from the last batch of the second epoch -> Option A

4. Quick Check:

   Printed loss = last batch loss [OK]

Hint: Printed loss after loop is last batch's loss [OK]

Common Mistakes:

• Thinking loss is summed or averaged automatically
• Assuming loss is from first batch
• Confusing loss tensor with scalar summary

4. You wrote this custom training loop but get an error:

for epoch in range(3):
    for x, y in dataset:
        with tf.GradientTape() as tape:
            predictions = model(x)
            loss = loss_fn(y, predictions)
        gradients = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

What is the likely cause of the error?

medium

A. Missing optimizer initialization

B. Using model.variables instead of model.trainable_variables in gradient calculation

C. Not calling loss.numpy() before gradient calculation

D. Forgetting to reset gradients after each batch

Solution

1. Step 1: Check gradient calculation variables

   Gradients must be computed with respect to trainable variables, not all variables.

2. Step 2: Identify the error source

   Using model.variables includes non-trainable variables, causing mismatch in optimizer step.

3. Final Answer:

   Using model.variables instead of model.trainable_variables in gradient calculation -> Option B

4. Quick Check:

   Gradient variables must be trainable variables [OK]

Hint: Use model.trainable_variables for gradients, not model.variables [OK]

Common Mistakes:

• Using all variables instead of trainable ones
• Forgetting to initialize optimizer
• Trying to convert loss to numpy before gradient step

5. You want to implement a custom training loop that stops early if validation loss does not improve for 3 epochs. Which approach gives you full control to do this?

hard

A. Write a for loop over epochs, compute training and validation loss each epoch, track best validation loss, and break loop if no improvement for 3 epochs

B. Use model.fit() with early stopping callback

C. Train with model.evaluate() and stop manually

D. Use a while loop without tracking validation loss

Solution

1. Step 1: Understand early stopping logic

   You must track validation loss each epoch and count epochs without improvement.

2. Step 2: Implement control flow in custom loop

   Using a for loop with manual checks and break statements gives full control to stop early.

3. Final Answer:

   Write a for loop over epochs, compute training and validation loss each epoch, track best validation loss, and break loop if no improvement for 3 epochs -> Option A

4. Quick Check:

   Manual loop + validation tracking = full control [OK]

Hint: Use manual loop with validation checks and break for early stopping [OK]

Common Mistakes:

• Relying only on model.fit() without custom logic
• Not tracking validation loss properly
• Using infinite loops without break conditions