Custom training step in TensorFlow - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output of a simple custom training step

What will be the printed loss value after running this custom training step once on a batch of data?

TensorFlow

import tensorflow as tf

class SimpleModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense = tf.keras.layers.Dense(1)

    def call(self, inputs):
        return self.dense(inputs)

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

model = SimpleModel()
optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

inputs = tf.constant([[1.0], [2.0], [3.0]])
targets = tf.constant([[2.0], [4.0], [6.0]])

loss_value = train_step(model, inputs, targets, optimizer)
print(float(loss_value))

A 0.0

B 4.6666665

C 1.0

D 9.0

💡 Hint

Think about the initial random weights and how the mean squared error is calculated.

The model starts with random weights, so predictions are initially far from targets. The mean squared error loss is calculated as the average of squared differences. Given the inputs and targets, the initial loss is approximately 4.6666665.

❓ Model Choice

intermediate

Choosing the correct custom training step for classification

Which custom training step correctly implements a single training iteration for a binary classification model using sigmoid activation and binary cross-entropy loss?

A

def train_step(model, inputs, targets, optimizer):
    with tf.GradientTape() as tape:
        logits = model(inputs)
        loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=targets, logits=logits))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

B

def train_step(model, inputs, targets, optimizer):
    with tf.GradientTape() as tape:
        predictions = tf.sigmoid(model(inputs))
        loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(targets, predictions))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

C

def train_step(model, inputs, targets, optimizer):
    with tf.GradientTape() as tape:
        logits = model(inputs)
        loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(targets, tf.sigmoid(logits)))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

D

def train_step(model, inputs, targets, optimizer):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=targets, logits=predictions))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

💡 Hint

Remember that sigmoid_cross_entropy_with_logits expects raw logits, not probabilities.

Option A correctly uses raw logits from the model and applies tf.nn.sigmoid_cross_entropy_with_logits, which expects logits and labels. Option A and C apply sigmoid before loss, which is incorrect for this function. Option A uses predictions as logits but predictions are already sigmoid outputs, causing wrong loss calculation.

❓ Hyperparameter

advanced

Effect of learning rate in a custom training step

If you increase the learning rate too much in a custom training step using SGD optimizer, what is the most likely effect on training?

A The model ignores the learning rate and trains normally

B The model's accuracy improves but loss remains constant

C The model's loss oscillates or diverges and fails to converge

D The model converges faster and reaches a lower loss without issues

💡 Hint

Think about what happens when steps are too large in gradient descent.

A very high learning rate causes the optimizer to take too large steps, overshooting minima. This leads to oscillations or divergence in loss, preventing convergence.

🔧 Debug

advanced

Identify the error in this custom training step

What error will this custom training step code raise when executed?

TensorFlow

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

A ValueError: No gradients provided for any variable

B TypeError: tape.gradient() missing required positional argument 'sources'

C TypeError: zip argument #1 must support iteration

D No error, runs successfully

💡 Hint

Check the arguments passed to tape.gradient and what it returns.

The tape.gradient call is missing the variables to compute gradients with respect to, so it returns None. Passing None to zip causes optimizer.apply_gradients to raise ValueError about no gradients provided.

🧠 Conceptual

expert

Why use a custom training step instead of model.fit?

Which is the best reason to implement a custom training step in TensorFlow instead of using the built-in model.fit method?

A Because custom training steps automatically improve model accuracy

B Because model.fit does not support GPU acceleration

C To avoid writing any code for backpropagation and optimization

D To have full control over the training loop, including custom loss functions, metrics, and gradient updates

💡 Hint

Think about flexibility and customization needs in training.

Custom training steps let you control every detail of training, such as applying different losses, multiple optimizers, or complex logic. model.fit is simpler but less flexible. model.fit supports GPU and backpropagation internally. Custom steps do not guarantee accuracy improvements automatically.

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