Understanding Linear Regression with TensorFlow in Python

Rate this post

πŸ’‘ Problem Formulation: Understanding how to implement linear regression models is essential for both novice and veteran data scientists. In this article, we explore how the popular machine-learning library TensorFlow assists with building such models in Python. Whether the task is to predict housing prices or to estimate a trend line for statistical data, your input consists of independent variables (features) and dependent variables (target values), and the output is a model defining the relationship between these variables.

Method 1: Using TensorFlow’s High-Level Keras API

TensorFlow houses a high-level API called Keras, which provides a straightforward way to create models, including linear regression. This approach involves specifying layers and compiling the model with a loss function and an optimizer.

Here’s an example:

import tensorflow as tf

# Model creation using the Sequential API
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

# Model compilation with mean squared error loss and a simple optimizer
model.compile(optimizer='sgd', loss='mean_squared_error')

# Fake data for training
x_train = [1, 2, 3, 4]
y_train = [2, 4, 6, 8]

# Training the model
model.fit(x_train, y_train, epochs=50)

Output: ‘Epoch 50/50…‘ indicating the training process.

This code snippet shows the creation of a simple linear model, which includes one dense layer (the standard layer type for linear regression in Keras). The model predicts a single continuous output (hence Dense(1)) from a single input feature. The ‘sgd’ optimizer is a commonly used algorithm that represents Stochastic Gradient Descent.

Method 2: Custom Training Loops

Going lower-level, TensorFlow allows you to create custom training loops for linear regression, providing more control over the training process. This is done by explicitly defining the weight variables, the loss computation, and the optimization step.

Here’s an example:

import tensorflow as tf

# Variables for weights and bias
W = tf.Variable(.1, dtype=tf.float32)
b = tf.Variable(.1, dtype=tf.float32)

# Training data
X = tf.constant([1, 2, 3, 4], dtype=tf.float32)
Y = tf.constant([2, 4, 6, 8], dtype=tf.float32)

# Training loop
for i in range(100):
    with tf.GradientTape() as tape:
        # Predicted Y
        Y_hat = W * X + b
        # Mean squared error
        loss = tf.reduce_mean(tf.square(Y - Y_hat))
    # Calculate gradients
    gradients = tape.gradient(loss, [W, b])

    # Update W and b
    W.assign_sub(0.01 * gradients[0])
    b.assign_sub(0.01 * gradients[1])

print(W.numpy(), b.numpy())

Output: ‘(array close to 2.0), (array close to 0.0)‘ representing the learned weights and bias.

The code defines weight W and bias b as TensorFlow variables, which enable automatic differentiation. Through the tf.GradientTape() block, TensorFlow records operations for automatic differentiation. This allows the calculation of gradients, which are then used to adjust variables in the direction that minimizes the loss.

Method 3: Using TensorFlow Estimators

TensorFlow also provides Estimators, which are high-level representatives of complete models. The Linear Estimator is pre-built for linear regression tasks, making the process quite straightforward while affording scalability and easy production deployment.

Here’s an example:

import tensorflow as tf
from tensorflow import feature_column
from tensorflow.estimator import LinearRegressor

# Feature columns describe how to use the input
feature_columns = [feature_column.numeric_column("x", shape=[1])]

# Estimator for Linear Regression
estimator = LinearRegressor(feature_columns=feature_columns)

# Data in the form of "tf.data.Dataset"
train_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn(
    {"x": np.array([1, 2, 3, 4])}, np.array([2, 4, 6, 8]), shuffle=False, num_epochs=None)

# Train the Estimator
estimator.train(input_fn=train_input_fn, steps=250)

Output: ‘No explicit output but trained model parameters.

Estimators encapsulate training, evaluation, prediction, and export for serving within a few lines of code. The LinearRegressor is specifically designed for linear regression, and expects input functions to be in a specific format provided by numpy_input_fn. They work well in distributed settings and are a solid choice for production-ready models.

Method 4: TensorFlow’s low-level operations

For educational purposes or for implementing very custom behavior, you can leverage TensorFlow’s low-level operations to perform linear regression. This involves math operations similar to the algebra conducted during manual calculations.

Here’s an example:

import tensorflow as tf

# Placeholder for training data
X = tf.placeholder(tf.float32)
Y = tf.placeholder(tf.float32)

# Model parameters
W = tf.Variable(np.random.randn(), name="weight")
b = tf.Variable(np.random.randn(), name="bias")

# Linear model
prediction = tf.add(tf.multiply(X, W), b)

# Loss calculation
loss = tf.reduce_sum(tf.pow(prediction-Y, 2)) / (2 * n_samples)

# Optimizer
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)

# Initialize variables and run session
with tf.Session() as sess:
    # Training steps...
    for step in range(training_epochs):
        sess.run(optimizer, feed_dict={X: x_train, Y: y_train})

# Query the trained values for W and b
trained_weight = sess.run(W)
trained_bias = sess.run(b)

Output: ‘trained_weight‘ and ‘trained_bias‘ with the values learned during the training.

This method directly leverages TensorFlow’s capability to build a computational graph composed of tensors and operations. Following initialization and during each training step within the session, operations are executed and the model is updated. While offering maximum flexibility, this method requires in-depth knowledge of TensorFlow’s mechanics and is more verbose.

Bonus One-Liner Method 5: Use TensorFlow’s Linear Model Feature

TensorFlow provides a simple linear model feature through the tf.linalg.lstsq function, which solves the Least Squares problem directly and can be used for implementing linear regression in a single line.

Here’s an example:

coefficients = tf.linalg.lstsq(tf.constant(x_train, dtype=tf.float32, shape=(4,1)),
                                    tf.constant(y_train, dtype=tf.float32, shape=(4,1)))

Output: Coefficients that estimate the relationship between ‘x_train’ and ‘y_train’.

By providing input data structured as tensors, the tf.linalg.lstsq function quickly calculates the best-fitting line for the given data points, yielding the coefficients of the linear regression. This method is extremely concise and useful when a simple solution on small to medium datasets is required.


  • Method 1: Keras API. Strengths: Simplistic high-level API, quick prototyping, many auxiliary functions. Weaknesses: Less flexible than lower-level APIs.
  • Method 2: Custom Training Loops. Strengths: Granular control over training, customizable for specific needs. Weaknesses: More complex code, higher chance of introducing errors.
  • Method 3: TensorFlow Estimators. Strengths: Easy to scale and deploy, encapsulated features for complete model operations. Weaknesses: Maybe challenging to customize for non-standard tasks.
  • Method 4: Low-level Operations. Strengths: Maximum flexibility, detailed understanding of how TensorFlow works. Weaknesses: Verbose, steep learning curve.
  • Method 5: Linear Model Feature. Strengths: Extremely concise, direct solution approach. Weaknesses: Limited flexibility and scalability.