TensorFlow Training a Model

Training a model is the process of showing your neural network many examples and letting it adjust its internal numbers (weights) until its predictions become accurate. This topic covers the complete training process in TensorFlow: how to prepare data for training, what happens inside each training step, how to track progress, and how to avoid the most common training mistakes.

The Learning Process: A School Exam Analogy

Imagine a student studying for an exam. The student reads a practice question (input data), writes an answer (prediction), checks it against the answer key (loss calculation), identifies what they got wrong (backpropagation), and studies those specific mistakes (weight update). After hundreds of practice questions, the student's answers become reliable. TensorFlow training works exactly this way.

The model.fit() Method

The model.fit() method handles the entire training loop. You call it once, and TensorFlow runs forward passes, computes loss, runs backpropagation, and updates weights automatically for every batch and every epoch.

history = model.fit(
    x_train,          # Training inputs
    y_train,          # Correct labels
    epochs=50,        # How many full passes through the training data
    batch_size=32,    # How many samples to process at once
    validation_data=(x_val, y_val),  # Data for measuring generalization
    shuffle=True      # Shuffle training data each epoch
)

Understanding Epochs and Batches

Dataset: 10,000 training samples
Batch size: 100
Batches per epoch: 10,000 ÷ 100 = 100 batches

Timeline for 3 Epochs:

Epoch 1:
  Batch 1 (samples 1–100)   → forward pass → loss → update weights
  Batch 2 (samples 101–200) → forward pass → loss → update weights
  ...
  Batch 100 (samples 9901–10000) → forward pass → loss → update weights
  → Print epoch 1 training loss and accuracy

Epoch 2:
  Shuffle data
  Batch 1 (different 100 samples) → ...
  ...

Epoch 3: Same pattern

Smaller batch sizes update weights more frequently but with noisier gradients. Larger batch sizes produce smoother gradient estimates but require more memory. Batch sizes of 32 or 64 work well in most situations.

Training vs. Validation Data

You always split your dataset into at least two parts:

  • Training set — the model learns from this data (typically 70–80% of your data)
  • Validation set — the model never trains on this; you use it to measure how well the model generalizes (typically 10–15%)
  • Test set — held back until the very end for final evaluation (typically 10–15%)
Full Dataset (1,000 samples)
│
├── Training Set (800 samples) — model adjusts weights here
├── Validation Set (100 samples) — monitored after each epoch
└── Test Set (100 samples) — evaluated only once, at the very end

Why Validation Data Matters

A model can memorize training data without actually learning patterns. If you only measure accuracy on training data, you cannot detect this problem. Validation accuracy reveals whether the model truly learned or just memorized.

Signal to watch for:

Healthy training:
  Epoch 1:  train_loss=0.90, val_loss=0.88
  Epoch 10: train_loss=0.45, val_loss=0.47
  Epoch 30: train_loss=0.22, val_loss=0.24
  → Both losses decrease together ✓

Overfitting (memorization):
  Epoch 1:  train_loss=0.90, val_loss=0.88
  Epoch 10: train_loss=0.30, val_loss=0.32
  Epoch 30: train_loss=0.05, val_loss=0.65 ← val_loss rising!
  → Model memorized training data, fails on new data ✗

The History Object

model.fit() returns a history object that stores training metrics for every epoch. Use it to plot training progress:

history = model.fit(x_train, y_train, epochs=50,
                    validation_data=(x_val, y_val))

# Access stored metrics
train_loss = history.history['loss']
val_loss   = history.history['val_loss']
train_acc  = history.history['accuracy']
val_acc    = history.history['val_accuracy']

# Each list has one value per epoch
print(f"Final training accuracy:   {train_acc[-1]:.3f}")
print(f"Final validation accuracy: {val_acc[-1]:.3f}")

Splitting Data in TensorFlow

import tensorflow as tf
import numpy as np

# Create fake data
X = np.random.random((1000, 20)).astype('float32')
y = np.random.randint(0, 2, size=1000)

# Manual split
split = 800
x_train, x_val = X[:split], X[split:]
y_train, y_val = y[:split], y[split:]

# Or use sklearn for a cleaner split
from sklearn.model_selection import train_test_split
x_train, x_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, random_state=42
)

Training with tf.data Datasets

For large datasets that do not fit in memory, use the tf.data.Dataset API to stream data from disk during training.

# Create training dataset from arrays
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(1000).batch(32).prefetch(
    tf.data.AUTOTUNE
)

val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(32)

# model.fit accepts tf.data.Dataset directly
model.fit(train_dataset, epochs=50, validation_data=val_dataset)

The prefetch(AUTOTUNE) call prepares the next batch on the CPU while the current batch trains on the GPU, eliminating wait time between batches.

Dealing With Class Imbalance

If your training data has many more examples of one class than another (for example, 950 healthy patients and only 50 sick patients), the model may learn to always predict "healthy" and achieve 95% accuracy while being completely useless for detecting disease.

Fix this by passing class weights to model.fit():

# Tell TensorFlow to penalize mistakes on the minority class more heavily
class_weight = {
    0: 1.0,   # healthy — normal weight
    1: 19.0   # sick — 19× more important (950/50 = 19)
}

model.fit(x_train, y_train, epochs=50, class_weight=class_weight)

Monitoring Training With Verbose Settings

  • verbose=0 — silent; no output printed (useful in production scripts)
  • verbose=1 — shows a progress bar for each epoch (default)
  • verbose=2 — shows one line per epoch without the progress bar

Common Training Problems and Fixes

Loss Is NaN (Not a Number)

The learning rate is too large and the model diverged. Reduce the learning rate by 10× and restart training.

Loss Stops Decreasing Early

The model hit a local minimum or the learning rate is too small. Try a learning rate warmup schedule or a different optimizer.

Training Is Very Slow

Make sure TensorFlow uses your GPU. Also increase the batch size to make better use of parallel hardware, and use prefetch(AUTOTUNE) in your data pipeline.

Validation Loss Much Higher Than Training Loss

The model is overfitting. Add Dropout layers, reduce model capacity (fewer layers or neurons), use data augmentation, or collect more training data.

A Complete Training Example

import tensorflow as tf
import numpy as np

# Data
x = np.random.random((2000, 30)).astype('float32')
y = (x[:, 0] + x[:, 1] > 1).astype(int)  # Simple rule to generate labels
x_train, x_val = x[:1600], x[1600:]
y_train, y_val = y[:1600], y[1600:]

# Model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(30,)),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Train
history = model.fit(
    x_train, y_train,
    epochs=40,
    batch_size=32,
    validation_data=(x_val, y_val),
    verbose=2
)

# Evaluate
final_val_acc = history.history['val_accuracy'][-1]
print(f"\nFinal Validation Accuracy: {final_val_acc:.2%}")

This complete workflow — prepare data, define model, compile, train, evaluate — is the foundation of every TensorFlow project. The next topic shows you how to formally evaluate a trained model on the test set and interpret what the metrics mean for your specific problem.

Previous lesson
Back to course
Next lesson

Leave a Comment

Your email address will not be published. Required fields are marked *