Custom training step in TensorFlow - Interactive Code Practice

Practice - 5 Tasks

1fill in blank

easy

Complete the code to define a custom training step function that calculates the loss.

TensorFlow

def train_step(model, inputs, targets, loss_fn):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = [1](targets, predictions)
    return loss

Drag options to blanks, or click blank then click option'

Aloss_fn

Bmodel.compile

Ctf.reduce_sum

Dtf.Variable

Attempts:

3 left

💡 Hint

Common Mistakes

Using model.compile instead of the loss function.

Using tf.reduce_sum which does not compute loss directly.

The loss function is called with targets and predictions to compute the loss inside the training step.

2fill in blank

medium

Complete the code to compute gradients of the loss with respect to the model's trainable variables.

TensorFlow

def train_step(model, inputs, targets, loss_fn, optimizer):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_fn(targets, predictions)
    gradients = tape.[1](loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

Drag options to blanks, or click blank then click option'

Agradient_of

Bgradient_tape

Cgradient

Dgradients

Attempts:

3 left

💡 Hint

Common Mistakes

Using plural 'gradients' instead of 'gradient'.

Using non-existent methods like 'gradient_of'.

The method to compute gradients is called gradient on the tape object.

3fill in blank

hard

Fix the error in the custom training step by completing the code to update the model weights.

TensorFlow

def train_step(model, inputs, targets, loss_fn, optimizer):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_fn(targets, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    [1].apply_gradients(zip(gradients, model.trainable_variables))
    return loss

Drag options to blanks, or click blank then click option'

Aoptimizer

Bloss

Cmodel

Dtape

Attempts:

3 left

💡 Hint

Common Mistakes

Using loss or model instead of optimizer to apply gradients.

Trying to call apply_gradients on tape.

The optimizer object is used to apply gradients to update the model weights.

4fill in blank

hard

Fill both blanks to complete the custom training step that returns loss and accuracy.

TensorFlow

def train_step(model, inputs, targets, loss_fn, optimizer, metric):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_fn(targets, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    metric.[1](targets, predictions)
    return loss, metric.[2]()

Drag options to blanks, or click blank then click option'

Aupdate_state

Bresult

Creset_states

Dcompute

Attempts:

3 left

💡 Hint

Common Mistakes

Using reset_states instead of update_state to update metric.

Using compute which is not a TensorFlow metric method.

The metric's update_state method updates the metric with new data, and result returns the current metric value.

5fill in blank

hard

Fill all three blanks to complete a custom training loop that iterates over batches and updates metrics.

TensorFlow

for batch, (inputs, targets) in enumerate(dataset):
    loss = train_step(model, inputs, targets, loss_fn, optimizer)
    train_loss.[1](loss)
    train_accuracy.[2](targets, model(inputs, training=True))
    if batch % 10 == 0:
        print(f"Batch {batch}: Loss = {train_loss.[3]():.4f}, Accuracy = {train_accuracy.result():.4f}")

Drag options to blanks, or click blank then click option'

Aupdate_state

Bupdate

Cresult

Dreset_states

Attempts:

3 left

💡 Hint

Common Mistakes

Using update instead of update_state which is the correct method.

Using reset_states instead of update_state.

Using result() where update_state is needed.

update_state updates the metric with new values, and result returns the current metric value for printing.

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