TensorFlow
Linear Regression
Machine Learning
Deep Learning
Python

Tensorflow on simple linear regression

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

Simple linear regression is one of the easiest ways to understand how TensorFlow trains a model. The goal is to learn a straight-line relationship of the form y = mx + b, and TensorFlow handles the gradient calculations and optimization for you.

The Regression Problem

In simple linear regression, there is one input feature and one target value. The model learns two parameters:

  • 'm, the slope'
  • 'b, the intercept'

Given data points, training adjusts those parameters so the predicted values are close to the observed targets.

A small synthetic dataset looks like this:

python
1import numpy as np
2
3x = np.array([1, 2, 3, 4, 5], dtype=np.float32)
4y = np.array([3, 5, 7, 9, 11], dtype=np.float32)

This roughly follows y = 2x + 1.

Build It with Keras

The easiest TensorFlow solution is a one-layer Keras model.

python
1import tensorflow as tf
2import numpy as np
3
4x = np.array([1, 2, 3, 4, 5], dtype=np.float32)
5y = np.array([3, 5, 7, 9, 11], dtype=np.float32)
6
7model = tf.keras.Sequential([
8    tf.keras.layers.Input(shape=(1,)),
9    tf.keras.layers.Dense(1)
10])
11
12model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=0.01),
13              loss="mse")
14
15model.fit(x, y, epochs=500, verbose=0)
16
17print(model.predict(np.array([6], dtype=np.float32), verbose=0))

The Dense(1) layer is enough because a linear neuron with one input already represents mx + b.

Inspect the Learned Parameters

After training, you can inspect the learned slope and intercept.

python
weights, bias = model.layers[0].get_weights()
print("slope:", weights[0][0])
print("bias:", bias[0])

On this dataset, those values should end up close to 2 and 1.

This is one of the nicest things about starting with linear regression: the model is simple enough that the parameters are directly interpretable.

The Loss Function and Optimizer

The standard loss for linear regression is mean squared error.

python
loss = "mse"

This penalizes larger prediction errors more heavily than smaller ones.

The optimizer updates the parameters using gradients. In this example, stochastic gradient descent is enough:

python
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

TensorFlow computes the gradients automatically during training, so you do not have to derive update formulas by hand.

Manual TensorFlow Version

If you want to see the mechanics more directly, you can train the same model with tf.Variable and GradientTape.

python
1import tensorflow as tf
2import numpy as np
3
4x = tf.constant([1, 2, 3, 4, 5], dtype=tf.float32)
5y = tf.constant([3, 5, 7, 9, 11], dtype=tf.float32)
6
7m = tf.Variable(0.0)
8b = tf.Variable(0.0)
9optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
10
11for epoch in range(500):
12    with tf.GradientTape() as tape:
13        y_pred = m * x + b
14        loss = tf.reduce_mean(tf.square(y - y_pred))
15
16    grads = tape.gradient(loss, [m, b])
17    optimizer.apply_gradients(zip(grads, [m, b]))
18
19print(m.numpy(), b.numpy())

This version makes the training loop explicit and is useful for understanding what Keras is automating for you.

When Linear Regression Is a Good Baseline

Even in larger machine learning projects, a linear model is a useful baseline because:

  • it trains quickly
  • it is interpretable
  • it tells you whether a simple linear relationship already explains much of the data

If a linear model performs surprisingly well, it may save you from building a more complex system too early.

Common Pitfalls

  • Feeding inputs with the wrong shape when Keras expects one feature per sample.
  • Using too large a learning rate and causing training to diverge.
  • Expecting a linear model to fit nonlinear data well.
  • Treating TensorFlow as if it were only for deep networks when simple models are also valid.
  • Ignoring the learned weights instead of checking whether they make sense for the problem.

Summary

  • Simple linear regression in TensorFlow is just a one-neuron linear model.
  • Keras makes it easy with Dense(1) and mean squared error loss.
  • You can inspect the learned slope and intercept after training.
  • A manual GradientTape loop helps explain what TensorFlow is doing internally.
  • Linear regression is a strong baseline because it is fast, interpretable, and easy to debug.

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.