Gradient Descent by Hand in TensorFlow

The Story — Walking Down a Mountain in the Fog

You now know how to get gradients with tf.GradientTape. In this lesson you will use them to train. No Keras fit(), no magic. Just a loop you write yourself. Every deep learning library does exactly this inside.

The One-Line Update Rule

Warm-up: Find the Bottom of (w − 3)²

The lowest point is at w = 3. Let us start far away, at w = −2, and let gradient descent find it.

import tensorflow as tf

w = tf.Variable(-2.0)
lr = 0.2

for step in range(8):
    with tf.GradientTape() as tape:
        loss = (w - 3.0) ** 2
    grad = tape.gradient(loss, w)
    w.assign_sub(lr * grad)
    print(f"step {step}: loss={loss.numpy():7.3f}  grad={grad.numpy():+7.3f}  new w={w.numpy():.4f}")

Each step closes 40% of the remaining gap (because 1 − 2 × lr = 0.6). The steps get smaller as the slope gets flatter. That is natural: near the bottom, the gradient is small.

Try It — Learning Rate Playground

Real Training: Fit a Line From Scratch

Now a real model. We make data from the line y = 2x + 1 plus some noise. The model does not know 2 or 1. It must learn them.

tf.random.set_seed(42)
X = tf.random.uniform([100], -3, 3)
y = 2.0 * X + 1.0 + tf.random.normal([100], stddev=0.5)   # true w=2, b=1

w = tf.Variable(0.0)
b = tf.Variable(0.0)
lr = 0.05

for epoch in range(101):
    with tf.GradientTape() as tape:
        pred = w * X + b
        loss = tf.reduce_mean((pred - y) ** 2)       # mean squared error
    dw, db = tape.gradient(loss, [w, b])
    w.assign_sub(lr * dw)
    b.assign_sub(lr * db)
    if epoch % 20 == 0:
        print(f"epoch {epoch:3d}  loss={loss.numpy():.4f}  w={w.numpy():.3f}  b={b.numpy():.3f}")

Try It — Train a Line Live

The Same Loop With an Optimizer

Writing assign_sub for every variable gets boring. An optimizer does the update for you. SGD is exactly our hand-written rule. Later you will swap in Adam with one line.

w = tf.Variable(0.0)
b = tf.Variable(0.0)
opt = tf.keras.optimizers.SGD(learning_rate=0.05)

for epoch in range(101):
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((w * X + b - y) ** 2)
    grads = tape.gradient(loss, [w, b])
    opt.apply_gradients(zip(grads, [w, b]))        # w -= lr * dw ; b -= lr * db

print(f"SGD optimizer -> w={w.numpy():.3f}  b={b.numpy():.3f}  loss={loss.numpy():.4f}")

The Most Common Bug

This loop looks right. But on the second step it crashes. Can you see why?

w = tf.Variable(0.0)
for step in range(2):
    with tf.GradientTape() as tape:
        loss = (w - 3.0) ** 2
    grad = tape.gradient(loss, w)
    print(f"step {step}: type(w)={type(w).__name__:16s} grad={grad}")
    try:
        w = w - 0.1 * grad            # BUG: makes w a plain Tensor
    except TypeError as e:
        print("TypeError:", e)

Golden Rules