TensorFlow
Machine Learning
Deep Learning
AI Training
Neural Networks

TensorFlow Training

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Training in TensorFlow means more than calling fit on a model. A solid training setup includes an input pipeline, a model whose output matches the target format, a loss that matches the output, and metrics that tell you whether optimization is actually improving the behavior you care about.

Start with a Reliable Input Pipeline

TensorFlow training is easiest to scale when data enters through tf.data. Even a small example benefits from batching and prefetching because the same pattern still works when the dataset grows.

python
1import tensorflow as tf
2
3features = tf.constant([
4    [0.0, 0.0],
5    [0.0, 1.0],
6    [1.0, 0.0],
7    [1.0, 1.0],
8], dtype=tf.float32)
9
10labels = tf.constant([[0.0], [1.0], [1.0], [1.0]], dtype=tf.float32)
11
12dataset = tf.data.Dataset.from_tensor_slices((features, labels))
13dataset = dataset.shuffle(buffer_size=4).batch(2).prefetch(tf.data.AUTOTUNE)

This example is tiny, but the structure is already correct: tensors become a dataset, the dataset is batched, and the runtime can overlap input work with training.

Use model.fit for the Normal Case

For most supervised learning tasks, Keras model.fit is the right default. It handles gradient computation, optimizer steps, metric aggregation, callbacks, and validation without forcing you to write the training loop yourself.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(2,)),
5    tf.keras.layers.Dense(8, activation="relu"),
6    tf.keras.layers.Dense(1, activation="sigmoid")
7])
8
9model.compile(
10    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
11    loss=tf.keras.losses.BinaryCrossentropy(),
12    metrics=[tf.keras.metrics.BinaryAccuracy()]
13)
14
15history = model.fit(dataset, epochs=20, verbose=0)
16print(history.history["loss"][-1])

That is enough for many production models. The usual mistake is abandoning fit too early when there is no real need for a custom loop.

Switch to GradientTape Only for Custom Behavior

Custom training loops are useful when you need unusual loss composition, manual gradient accumulation, or tightly controlled reinforcement and sequence objectives. TensorFlow exposes that level through tf.GradientTape.

python
1import tensorflow as tf
2
3optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
4loss_fn = tf.keras.losses.BinaryCrossentropy()
5
6for epoch in range(3):
7    for batch_x, batch_y in dataset:
8        with tf.GradientTape() as tape:
9            predictions = model(batch_x, training=True)
10            loss = loss_fn(batch_y, predictions)
11
12        gradients = tape.gradient(loss, model.trainable_variables)
13        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
14
15    print(f"epoch={epoch + 1} loss={loss.numpy():.4f}")

The advantage is control. The cost is that you now own metric tracking, validation loops, checkpoint timing, and error handling yourself.

Monitor the Right Things

Loss is necessary but not always sufficient. Classification projects may care more about precision, recall, or area under the curve. Regression may care more about mean absolute error than mean squared error. The training setup should reflect the deployment goal instead of only reporting whichever metric is shortest to configure.

Validation also matters. A steadily falling training loss is not enough if validation performance stalls or degrades. In practice, callbacks such as early stopping and model checkpointing are often as important as the optimizer choice.

Reproducibility matters too. If you change seeds, batch order, optimizer settings, and model depth at the same time, the run history becomes hard to interpret. Training improves faster when experiments are small and logged clearly.

Common Pitfalls

  • Using an output activation and a loss function that disagree about whether predictions are logits or probabilities.
  • Feeding data through slow Python loops instead of a batched tf.data pipeline.
  • Writing a custom training loop when model.fit would already handle the problem cleanly.
  • Watching only training loss and missing overfitting on validation data.
  • Changing several training variables at once and then not knowing which change actually helped.

Summary

  • Good TensorFlow training starts with a clean input pipeline and a model-loss pairing that matches the task.
  • 'model.fit is the best default for standard supervised training.'
  • 'tf.GradientTape is valuable when you need custom optimization behavior.'
  • Metrics should reflect the actual business or modeling goal, not only the default loss.
  • Validation, checkpointing, and controlled experimentation matter as much as the model definition itself.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.