What is Custom training step in TensorFlow?

TensorFlowml~5 mins

Custom training step in TensorFlow

Choose your learning style10 modes available

Learn Why Deep Model Try Challenge Experiment Recall Metrics

Introduction

Sometimes, you want to control exactly how your model learns. A custom training step lets you do that by writing your own code for one round of learning.

You want to add special rules during training, like extra checks or custom loss calculations.

You need to update parts of the model differently, not just all at once.

You want to track extra information during training, like custom metrics.

You want to experiment with new ways to train models that built-in methods don't support.

Syntax

TensorFlow

def train_step(data):
    with tf.GradientTape() as tape:
        predictions = model(data[0], training=True)
        loss = loss_function(data[1], predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss.update_state(loss)
    train_accuracy.update_state(data[1], predictions)
    return {"loss": train_loss.result(), "accuracy": train_accuracy.result()}

The train_step function takes a batch of data and runs one learning round.

Use tf.GradientTape() to record operations for automatic differentiation.

Examples

Simple custom step returning only the loss.

TensorFlow

def train_step(data):
    x, y = data
    with tf.GradientTape() as tape:
        y_pred = model(x, training=True)
        loss = loss_fn(y, y_pred)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

Custom step that updates and returns loss and accuracy metrics.

TensorFlow

def train_step(data):
    x, y = data
    with tf.GradientTape() as tape:
        y_pred = model(x, training=True)
        loss = loss_fn(y, y_pred)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    train_loss_metric.update_state(loss)
    train_acc_metric.update_state(y, y_pred)
    return {"loss": train_loss_metric.result(), "accuracy": train_acc_metric.result()}

Sample Model

This code loads a small part of MNIST, builds a simple CNN, and runs one epoch using a custom training step. It prints the loss and accuracy after training.

TensorFlow

import tensorflow as tf

# Prepare simple data
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train[:1000].astype('float32') / 255.0
x_train = x_train[..., tf.newaxis]
y_train = y_train[:1000]

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(32)

# Build a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, 3, activation='relu'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10)
])

loss_fn = 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')

@tf.function
def train_step(data):
    x, y = data
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss.update_state(loss)
    train_accuracy.update_state(y, predictions)
    return {"loss": train_loss.result(), "accuracy": train_accuracy.result()}

# Run one epoch
for batch in train_ds:
    metrics = train_step(batch)

print(f"Loss: {metrics['loss']:.4f}, Accuracy: {metrics['accuracy']:.4f}")

OutputSuccess

Important Notes

Use @tf.function to speed up the training step by compiling it.

Reset metrics if you run multiple epochs to avoid mixing results.

Custom training steps give full control but require careful coding to avoid mistakes.

Summary

Custom training steps let you write exactly how your model learns each batch.

They use tf.GradientTape to calculate gradients and update weights.

You can track and return any metrics you want during training.

Practice

(1/5)

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

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:

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

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

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()

Recommended

Ready to prove it?

Pass the test and this pattern moves unknown strong while earning +1 Grip toward your certificate.

10 questions~5 mineasy hard

Custom training step in TensorFlow

Practice

Solution

Step 1: Understand the role of a training step

Step 2: Identify what custom training step changes

Final Answer:

Quick Check:

Solution

Step 1: Recall TensorFlow's gradient recording method

Step 2: Match correct syntax for gradient recording

Final Answer:

Quick Check:

Solution

Step 1: Calculate initial predictions and loss

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

Final Answer:

Quick Check:

Solution

Step 1: Check tape.gradient usage

Step 2: Identify missing argument

Final Answer:

Quick Check:

Solution

Step 1: Understand metric update methods

Step 2: Correctly update and return metric results

Final Answer:

Quick Check:

Ready to prove it?