Mahmoud Abas, Full-Stack Software Engineer & AI

Neural Networks - ML Review and Gradient Descent Example

In the print dialog, choose "Save as PDF" as the destination.

Neural Networks - Lesson 3

ML Review and Gradient Descent Example

Building and training a linear regression model with TensorFlow Core APIs

The ML Process

Solve linear regression via gradient descent, following the standard ML workflow

Load the dataset
Build a multiple linear regression model
Evaluate the performance of the model

python

import tensorflow as tf
import pandas as pd
from matplotlib import pyplot as plt
tf.random.set_seed(22)

Imports and a fixed random seed, so the run is reproducible.

Generating Data from a Known Distribution

Y = b + X₁ w₁ + X₂ w₂

X_1 and X_2 are uniform on [0, 10]
Known ground truth: b = 1.5, w_1 = 2, w_2 = 5 - the goal is to recover these values via training

python

num_obs = 150
x1 = tf.random.uniform([num_obs], minval=0, maxval=10, dtype=tf.float32)
x2 = tf.random.uniform([num_obs], minval=0, maxval=10, dtype=tf.float32)

b, w_1, w_2 = 1.5, 2.0, 5.0
y = b + w_1 * x1 + w_2 * x2

dataset = tf.stack([y, x1, x2], axis=1)
# shape (150, 3)

dataset stacks y, x1, x2 into one matrix, shape (150, 3).

Visualizing the Data

A 3D scatter shows how y depends jointly on x1 and x2
2D scatter plots show each feature's individual relationship with y

python

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(x1.numpy(), x2.numpy(), y.numpy(), c='red', marker='o')
ax.set_xlabel('x1'); ax.set_ylabel('x2'); ax.set_zlabel('y')
plt.title("3D Scatter: (x1, x2, y)")

Train/Test Split

Shuffle before splitting, to avoid biased splits (for example, data ordered by some hidden pattern)
125 training observations, 25 test observations, a ~83/17 split

python

dataset_shuffled = tf.random.shuffle(dataset, seed=22)
train_data, test_data = dataset_shuffled[:125], dataset_shuffled[125:]

x_train, y_train = train_data[:, 1:], train_data[:, 0]
x_test, y_test = test_data[:, 1:], test_data[:, 0]
# x_train shape: (125, 2)

x_train shape: (125, 2).

Building the Linear Regression Model

Y = Xw + b

Implemented as a tf.Module subclass; weights are lazily initialized on first call
@tf.function traces the call into a graph, for performance and for saving/loading

python

class LinearRegression(tf.Module):
  def __init__(self):
    self.built = False

  @tf.function
  def __call__(self, x):
    if not self.built:
      rand_w = tf.random.uniform(shape=[x.shape[-1], 1])
      rand_b = tf.random.uniform(shape=[])
      self.w = tf.Variable(rand_w)
      self.b = tf.Variable(rand_b)
      self.built = True
    y = tf.add(tf.matmul(x, self.w), self.b)
    return tf.squeeze(y, axis=1)

The Loss Function: Mean Squared Error

Mean Squared Error (MSE): The standard loss for regression problems: it measures how far the predictions are from the true targets. Training's goal is to find w and b that minimize it.

MSE = 1m∑mi=1(ŷ_i - y_i)²

python

def mse_loss(y_pred, y):
  return tf.reduce_mean(tf.square(y_pred - y))

Mini-Batches with tf.data.Dataset

Mini-batch training balances memory efficiency and convergence speed
tf.data.Dataset handles shuffling and batching the input pipeline

python

batch_size = 32
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=x_train.shape[0]).batch(batch_size)

test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_dataset = test_dataset.shuffle(buffer_size=x_test.shape[0]).batch(batch_size)
# Batches: (32,2), (32,2), (32,2), (29,2)

Effect of Batch Size

Larger batch size: faster training, better hardware utilization, but more memory and can hurt generalization
Smaller batch size: better generalization, less memory, but slower and noisier gradients

Batch Size	When to Use
1	Reinforcement learning, online learning
8-16	Small datasets, memory-constrained GPUs
32 (default)	Good balance of speed & generalization
64-128	Large datasets, powerful GPUs
256+	Very large datasets, TPU/GPU clusters

Effect of Learning Rate

The learning rate controls how much weights update at each step of gradient descent
Too high: risk of overshooting, loss may diverge. Too low: very slow convergence, may stall

Model Type	Suggested Learning Rate
Simple Linear Regression	0.01 - 0.1
Deep Neural Networks	0.001 - 0.01
CNNs	0.0001 - 0.01
Transformers (e.g. BERT)	0.00001 - 0.0001

The Training Loop (Gradient Descent)

Gradient Descent: An iterative method that updates parameters by taking a step in the opposite direction of their computed gradient, one step down the loss surface at a time.

Predict

Compute loss & gradients

Update parameters

Each step: predict, compute loss and gradients, then update parameters against the gradient
v.assign_sub(learning_rate * g) moves each variable one step down the loss surface

python

epochs, learning_rate = 100, 0.01
train_losses, test_losses = [], []

for epoch in range(epochs):
  batch_losses_train = []
  for x_batch, y_batch in train_dataset:
    with tf.GradientTape() as tape:
      y_pred_batch = lin_reg(x_batch)
      batch_loss = mse_loss(y_pred_batch, y_batch)
    grads = tape.gradient(batch_loss, lin_reg.variables)
    for g, v in zip(grads, lin_reg.variables):
      v.assign_sub(learning_rate * g)
    batch_losses_train.append(batch_loss)

w ≈ [2.03, 5.02], b ≈ 1.18

Result after 100 epochs: train loss 0.014, against the true values w = [2, 5], b = 1.5

Plotting Loss Over Time

Plotting train vs. test loss over epochs confirms the model converges and generalizes (no overfitting)

python

plt.plot(range(epochs), train_losses, label="Training loss")
plt.plot(range(epochs), test_losses, label="Testing loss")
plt.xlabel("Epoch")
plt.ylabel("Mean squared error loss")
plt.legend()
plt.title("MSE loss vs training iterations")

Key Takeaways

The ML workflow: load data, build a model, define a loss, train, evaluate on held-out data
Multiple linear regression can be implemented directly with TensorFlow Core APIs (tf.Module, tf.matmul)
MSE is the standard loss for regression; gradient descent minimizes it step by step
tf.GradientTape() computes gradients; v.assign_sub(lr * grad) applies the update
Batch size and learning rate are key hyperparameters trading off speed, stability, and generalization
Tracking train and test loss together reveals whether the model generalizes or overfits

Open this lesson

Mahmoud AbasML Review and Gradient Descent Example

num_obs = 150 x1 = tf.random.uniform([num_obs], minval=0, maxval=10, dtype=tf.float32) x2 = tf.random.uniform([num_obs], minval=0, maxval=10, dtype=tf.float32) b, w_1, w_2 = 1.5, 2.0, 5.0 y = b + w_1 * x1 + w_2 * x2 dataset = tf.stack([y, x1, x2], axis=1) # shape (150, 3)

fig = plt.figure(figsize=(8, 6)) ax = fig.add_subplot(111, projection='3d') ax.scatter(x1.numpy(), x2.numpy(), y.numpy(), c='red', marker='o') ax.set_xlabel('x1'); ax.set_ylabel('x2'); ax.set_zlabel('y') plt.title("3D Scatter: (x1, x2, y)")

dataset_shuffled = tf.random.shuffle(dataset, seed=22) train_data, test_data = dataset_shuffled[:125], dataset_shuffled[125:] x_train, y_train = train_data[:, 1:], train_data[:, 0] x_test, y_test = test_data[:, 1:], test_data[:, 0] # x_train shape: (125, 2)

class LinearRegression(tf.Module): def __init__(self): self.built = False @tf.function def __call__(self, x): if not self.built: rand_w = tf.random.uniform(shape=[x.shape[-1], 1]) rand_b = tf.random.uniform(shape=[]) self.w = tf.Variable(rand_w) self.b = tf.Variable(rand_b) self.built = True y = tf.add(tf.matmul(x, self.w), self.b) return tf.squeeze(y, axis=1)

batch_size = 32 train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)) train_dataset = train_dataset.shuffle(buffer_size=x_train.shape[0]).batch(batch_size) test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)) test_dataset = test_dataset.shuffle(buffer_size=x_test.shape[0]).batch(batch_size) # Batches: (32,2), (32,2), (32,2), (29,2)

Batch Size

When to Use

Reinforcement learning, online learning

8-16

Small datasets, memory-constrained GPUs

32 (default)

Good balance of speed & generalization

64-128

Large datasets, powerful GPUs

256+

Very large datasets, TPU/GPU clusters

Model Type

Suggested Learning Rate

Simple Linear Regression

0.01 - 0.1

Deep Neural Networks

0.001 - 0.01

CNNs

0.0001 - 0.01

Transformers (e.g. BERT)

0.00001 - 0.0001

epochs, learning_rate = 100, 0.01 train_losses, test_losses = [], [] for epoch in range(epochs): batch_losses_train = [] for x_batch, y_batch in train_dataset: with tf.GradientTape() as tape: y_pred_batch = lin_reg(x_batch) batch_loss = mse_loss(y_pred_batch, y_batch) grads = tape.gradient(batch_loss, lin_reg.variables) for g, v in zip(grads, lin_reg.variables): v.assign_sub(learning_rate * g) batch_losses_train.append(batch_loss)

plt.plot(range(epochs), train_losses, label="Training loss") plt.plot(range(epochs), test_losses, label="Testing loss") plt.xlabel("Epoch") plt.ylabel("Mean squared error loss") plt.legend() plt.title("MSE loss vs training iterations")