Custom training step in TensorFlow - Deep Dive

Overview - Custom training step

What is it?

A custom training step is a way to control exactly how a machine learning model learns from data by writing your own code for one cycle of training. Instead of using built-in training functions, you define how the model processes inputs, calculates errors, and updates itself. This lets you add special behaviors or experiments that standard training does not support.

Why it matters

Custom training steps exist because sometimes the default training methods are too simple or rigid for complex tasks. Without them, you cannot easily try new ideas like custom loss calculations, multiple updates per batch, or special data handling. This limits innovation and fine control in building smarter models.

Where it fits

Before learning custom training steps, you should understand basic TensorFlow model building and the standard training process using fit(). After mastering custom steps, you can explore advanced topics like distributed training, mixed precision, and custom optimizers.

Mental Model

Core Idea

A custom training step is writing your own instructions for how the model learns from one batch of data, replacing the automatic training loop.

Think of it like...

It's like cooking your own recipe instead of following a pre-made meal kit; you decide every ingredient and step to create a dish exactly how you want.

┌─────────────────────────────┐
│       Custom Training Step   │
├──────────────┬──────────────┤
│ Input batch  │ Data batch   │
│              │ (features &  │
│              │ labels)      │
├──────────────┴──────────────┤
│ Forward pass: model predicts │
├──────────────┬──────────────┤
│ Loss calc    │ Compare pred  │
│              │ to labels     │
├──────────────┴──────────────┤
│ Backprop: gradients computed │
├──────────────┬──────────────┤
│ Optimizer    │ Update weights│
│ step         │              │
└──────────────┴──────────────┘

Build-Up - 6 Steps

1

Foundation

Understanding standard training loop

Concept: Learn how TensorFlow's built-in fit() method trains models automatically.

TensorFlow's fit() method takes your data and runs many training cycles (epochs). Each cycle processes batches of data, calculates the error (loss), and updates the model weights to improve predictions. This is done behind the scenes, so you don't see the details.

Result

You get a trained model without writing code for each training step.

Knowing the default training process helps you understand what you control when writing a custom training step.

2

Foundation

Basics of forward and backward passes

Concept: Understand how models make predictions and learn from errors using gradients.

In each training step, the model first makes predictions from inputs (forward pass). Then, it compares predictions to true answers to calculate loss. Next, it computes gradients showing how to change weights to reduce loss (backward pass). Finally, it updates weights using an optimizer.

Result

You see the core math behind learning: prediction, error, gradient, update.

Grasping forward and backward passes is essential to customize training behavior.

3

Intermediate

Writing a simple custom training step

🤔Before reading on: do you think you need to manually compute gradients or does TensorFlow do it automatically? Commit to your answer.

Concept: Learn to write a function that runs one training step using TensorFlow's GradientTape to compute gradients.

You write a function that takes a batch of data, runs the model to get predictions, calculates loss, then uses tf.GradientTape to watch operations and compute gradients. Then you apply gradients with an optimizer to update weights. This replaces the automatic fit() step.

Result

You can run one training step manually and control every detail.

Understanding GradientTape unlocks the ability to customize exactly how learning happens.

4

Intermediate

Adding metrics and multiple losses

🤔Before reading on: can you track accuracy manually inside a custom training step or is it automatic? Commit to your answer.

Concept: Extend the custom step to calculate and update metrics like accuracy, and combine multiple loss functions.

Inside the training step, after computing loss, you can update metrics by calling metric.update_state(). You can also compute several losses (e.g., classification and regularization) and add them together before backpropagation. This gives fine control over what the model optimizes.

Result

You get detailed feedback on training progress and can optimize complex objectives.

Manual metric updates let you monitor training precisely and customize learning goals.

5

Advanced

Handling custom gradients and manual weight updates

🤔Before reading on: do you think you can modify gradients before applying them? Commit to your answer.

Concept: Learn to manipulate gradients before applying them, enabling techniques like gradient clipping or custom updates.

After computing gradients with GradientTape, you can change them (e.g., clip large values to prevent exploding gradients). Then you apply these modified gradients with the optimizer. You can also update weights manually by assigning new values, bypassing the optimizer if needed.

Result

You gain powerful control to improve training stability and experiment with new update rules.

Modifying gradients is key to advanced training tricks and avoiding common training failures.

6

Expert

Integrating custom training steps with distributed training

🤔Before reading on: do you think custom training steps work automatically with multiple GPUs or do you need extra code? Commit to your answer.

Concept: Understand how to adapt custom training steps to run efficiently on multiple devices using TensorFlow's distribution strategies.

When using multiple GPUs or TPUs, you wrap your custom training step inside a distribution strategy's run() method. You must aggregate gradients and losses across devices manually. This requires careful coding to ensure synchronization and correct updates.

Result

Your custom training step can scale to large datasets and hardware setups.

Knowing how to integrate with distribution strategies is essential for production-scale custom training.

Under the Hood

TensorFlow uses automatic differentiation to record operations during the forward pass inside a GradientTape context. It then computes gradients by applying the chain rule backward through the recorded operations. These gradients represent how much each weight affects the loss. The optimizer uses these gradients to update weights, typically by moving them slightly in the direction that reduces loss.

Why designed this way?

This design allows flexible, efficient gradient computation without manual math. Automatic differentiation supports complex models and custom operations. Writing custom training steps with GradientTape gives users full control while leveraging TensorFlow's optimized backend.

┌───────────────┐
│ Input batch   │
└──────┬────────┘
       │
┌──────▼───────┐
│ Forward pass │
│ (record ops) │
└──────┬───────┘
       │
┌──────▼───────┐
│ Loss compute │
└──────┬───────┘
       │
┌──────▼─────────────┐
│ Backward pass       │
│ (auto differentiation)│
└──────┬─────────────┘
       │
┌──────▼─────────────┐
│ Gradient update     │
│ (optimizer step)   │
└────────────────────┘

Myth Busters - 4 Common Misconceptions

Quick: Does writing a custom training step mean you must write all code from scratch, including optimizer logic? Commit yes or no.

Common Belief: People often think custom training steps require rewriting the optimizer and all training logic manually.

Reality: You only write the training step logic; you still use TensorFlow's built-in optimizers to apply gradients.

Why it matters: Believing you must rewrite optimizers scares beginners and wastes effort, slowing learning and experimentation.

Quick: Is it true that custom training steps always run slower than fit()? Commit yes or no.

Common Belief: Many assume custom training steps are slower because they are manual.

Reality: Custom steps can be as fast or faster if optimized well, and allow performance improvements like gradient accumulation or mixed precision.

Why it matters: Thinking custom steps are slow may prevent users from trying powerful customizations that improve training.

Quick: Do you think metrics like accuracy update automatically inside custom training steps? Commit yes or no.

Common Belief: Some believe metrics update automatically without explicit code in custom steps.

Reality: You must manually update metrics inside the custom training step for them to track correctly.

Why it matters: Missing metric updates leads to wrong or missing training feedback, causing confusion about model progress.

Quick: Is it safe to modify gradients arbitrarily before applying them? Commit yes or no.

Common Belief: People sometimes think any gradient modification is allowed without consequences.

Reality: Improper gradient changes can break training or cause divergence; modifications must be done carefully with understanding.

Why it matters: Misusing gradient modifications can silently ruin model training, wasting time and resources.

Expert Zone

1

Custom training steps allow mixing eager execution and graph mode for debugging and speed, a subtle balance experts exploit.

2

GradientTape can watch only specific variables or operations, enabling fine-grained control over what gradients are computed.

3

When using custom training steps with tf.function, tracing and retracing behavior affects performance and must be managed carefully.

When NOT to use

Avoid custom training steps if your training needs are simple and fit() covers them well, as custom steps add complexity. For very large-scale distributed training, consider high-level APIs like tf.distribute or specialized frameworks that handle synchronization automatically.

Production Patterns

In production, custom training steps are used to implement advanced techniques like adversarial training, multi-task learning, or custom regularization. They also enable integration with complex data pipelines and hardware accelerators, allowing fine-tuned performance and novel research experiments.

Connections

Automatic differentiation

Custom training steps build directly on automatic differentiation to compute gradients.

Understanding automatic differentiation clarifies how gradients are computed efficiently and correctly during custom training.

Control systems engineering

Both involve feedback loops where outputs influence inputs to improve system behavior.

Seeing training as a feedback control loop helps grasp why gradients guide weight updates to reduce errors.

Cooking recipes

Custom training steps are like writing your own recipe instead of following a fixed one.

This connection highlights the freedom and responsibility in customizing processes to achieve desired results.

Common Pitfalls

#1 Not using GradientTape to watch operations, so gradients are not computed.

Wrong approach:

def train_step(inputs, labels):
predictions = model(inputs)
loss = loss_fn(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables) # Error: tape not defined
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

Correct approach:

def train_step(inputs, labels):
with tf.GradientTape() as tape:
predictions = model(inputs)
loss = loss_fn(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

Root cause:

Forgetting to create a GradientTape context means TensorFlow cannot record operations to compute gradients.

#2 Not updating metrics inside the custom training step, so metrics stay at zero.

Wrong approach:

def train_step(inputs, labels):
with tf.GradientTape() as tape:
predictions = model(inputs)
loss = loss_fn(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
# Missing metric update

Correct approach:

def train_step(inputs, labels):
with tf.GradientTape() as tape:
predictions = model(inputs)
loss = loss_fn(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
accuracy_metric.update_state(labels, predictions)

Root cause:

Metrics require explicit calls to update_state; they do not update automatically.

#3 Applying gradients without checking for None values, causing errors.

Wrong approach:

gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables)) # May fail if gradients contain None

Correct approach:

gradients = tape.gradient(loss, model.trainable_variables)
grads_and_vars = [(g, v) for g, v in zip(gradients, model.trainable_variables) if g is not None]
optimizer.apply_gradients(grads_and_vars)

Root cause:

Some variables may not affect loss and have None gradients; applying them causes errors.

Key Takeaways

Custom training steps let you control exactly how a model learns from data, beyond built-in methods.

Using tf.GradientTape is essential to compute gradients automatically during your custom step.

You must manually update metrics and handle multiple losses inside custom training steps.

Advanced uses include modifying gradients and integrating with distributed training for scale.

Understanding the internal mechanics of training empowers you to innovate and debug effectively.

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