Introduction
Linear regression in TensorFlow involves two core data structures: the attribute matrix (X) containing input features and the target matrix (Y) containing the values to predict. Properly shaping and preparing these matrices is essential — TensorFlow operations expect specific dimensions, and mismatched shapes are the most common source of errors in regression models.
Understanding the Matrices
Attribute Matrix (X): Shape (n_samples, n_features). Each row is one observation, each column is a feature.
Target Matrix (Y): Shape (n_samples, 1) or (n_samples,). The output variable to predict.
1import numpy as np
2
3# Example: predicting house prices
4# Features: square_feet, bedrooms, age
5X = np.array([
6 [1500, 3, 10],
7 [2000, 4, 5],
8 [1200, 2, 20],
9 [1800, 3, 8],
10 [2500, 5, 2]
11], dtype=np.float32) # Shape: (5, 3)
12
13# Target: price in thousands
14Y = np.array([300, 400, 250, 350, 500], dtype=np.float32) # Shape: (5,)
15Y = Y.reshape(-1, 1) # Shape: (5, 1) — required for matrix operations
Method 1: TensorFlow Keras (Recommended)
The simplest approach using tf.keras:
1import tensorflow as tf
2import numpy as np
3from sklearn.model_selection import train_test_split
4from sklearn.preprocessing import StandardScaler
5
6# Generate sample data
7np.random.seed(42)
8X = np.random.randn(1000, 3).astype(np.float32)
9true_weights = np.array([[2.0], [3.0], [-1.0]])
10Y = X @ true_weights + 0.5 * np.random.randn(1000, 1).astype(np.float32)
11
12# Split data
13X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2)
14
15# Normalize features
16scaler = StandardScaler()
17X_train = scaler.fit_transform(X_train)
18X_test = scaler.transform(X_test)
19
20# Build model
21model = tf.keras.Sequential([
22 tf.keras.layers.Dense(1, input_shape=(3,)) # Linear: Y = X @ W + b
23])
24
25model.compile(optimizer='adam', loss='mse', metrics=['mae'])
26
27# Train
28history = model.fit(X_train, Y_train, epochs=100, batch_size=32,
29 validation_split=0.2, verbose=0)
30
31# Evaluate
32loss, mae = model.evaluate(X_test, Y_test)
33print(f"Test MSE: {loss:.4f}, MAE: {mae:.4f}")
34
35# Get learned weights
36weights, bias = model.layers[0].get_weights()
37print(f"Weights: {weights.flatten()}") # Close to [2.0, 3.0, -1.0]
38print(f"Bias: {bias[0]:.4f}") # Close to 0.0
Method 2: Low-Level TensorFlow
For understanding the math behind linear regression:
1import tensorflow as tf
2import numpy as np
3
4# Data
5X_data = np.random.randn(100, 2).astype(np.float32)
6Y_data = (X_data @ np.array([[3.0], [-2.0]]) + 1.0).astype(np.float32)
7
8# Convert to tensors
9X = tf.constant(X_data)
10Y = tf.constant(Y_data)
11
12# Variables (learnable parameters)
13W = tf.Variable(tf.random.normal([2, 1]))
14b = tf.Variable(tf.zeros([1]))
15
16# Training loop
17learning_rate = 0.1
18for epoch in range(200):
19 with tf.GradientTape() as tape:
20 predictions = tf.matmul(X, W) + b # Y_hat = X @ W + b
21 loss = tf.reduce_mean(tf.square(predictions - Y)) # MSE
22
23 gradients = tape.gradient(loss, [W, b])
24 W.assign_sub(learning_rate * gradients[0])
25 b.assign_sub(learning_rate * gradients[1])
26
27 if epoch % 50 == 0:
28 print(f"Epoch {epoch}: Loss = {loss.numpy():.4f}")
29
30print(f"Learned W: {W.numpy().flatten()}") # Close to [3.0, -2.0]
31print(f"Learned b: {b.numpy()[0]:.4f}") # Close to 1.0
For small datasets, solve analytically:
1import tensorflow as tf
2import numpy as np
3
4X_data = np.random.randn(100, 2).astype(np.float32)
5Y_data = (X_data @ np.array([[3.0], [-2.0]]) + 1.0).astype(np.float32)
6
7# Add bias column (column of ones)
8X_with_bias = tf.concat([tf.ones([100, 1]), tf.constant(X_data)], axis=1)
9
10# Normal equation: W = (X^T X)^(-1) X^T Y
11XtX = tf.matmul(tf.transpose(X_with_bias), X_with_bias)
12XtY = tf.matmul(tf.transpose(X_with_bias), tf.constant(Y_data))
13W = tf.matmul(tf.linalg.inv(XtX), XtY)
14
15print(f"Weights (bias, w1, w2): {W.numpy().flatten()}")
16# Close to [1.0, 3.0, -2.0]
Loading Data from CSV
1import pandas as pd
2import numpy as np
3import tensorflow as tf
4
5# Load CSV
6df = pd.read_csv('housing.csv')
7
8# Separate features and target
9feature_cols = ['square_feet', 'bedrooms', 'age']
10X = df[feature_cols].values.astype(np.float32)
11Y = df['price'].values.astype(np.float32).reshape(-1, 1)
12
13# Using tf.data for efficient batching
14dataset = tf.data.Dataset.from_tensor_slices((X, Y))
15dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)
16
17# Train with dataset
18model.fit(dataset, epochs=50)
Feature Normalization
Always normalize features before training — linear regression converges much faster with normalized inputs:
1from sklearn.preprocessing import StandardScaler
2
3scaler = StandardScaler()
4X_train_scaled = scaler.fit_transform(X_train) # Mean=0, Std=1
5X_test_scaled = scaler.transform(X_test) # Same transform
6
7# Or with TensorFlow's Normalization layer
8normalizer = tf.keras.layers.Normalization(axis=-1)
9normalizer.adapt(X_train)
10
11model = tf.keras.Sequential([
12 normalizer,
13 tf.keras.layers.Dense(1)
14])
Common Pitfalls
Shape mismatch: X must be (n_samples, n_features) and Y must be (n_samples, 1). A common error is Y being (n_samples,) — use Y.reshape(-1, 1).
Feature scaling: Without normalization, features with large ranges dominate the loss. Always standardize (zero mean, unit variance) or normalize (0-1 range) before training.
dtype mismatch: TensorFlow defaults to float32. Ensure X and Y are both float32, not float64 (NumPy's default).
Learning rate too high: For gradient descent, a high learning rate causes divergence (loss goes to infinity). Start with 0.001 and adjust.
Overfitting with many features: Linear regression with many features relative to samples can overfit. Use L1/L2 regularization: tf.keras.layers.Dense(1, kernel_regularizer=tf.keras.regularizers.l2(0.01)).
Summary
The attribute matrix X has shape (n_samples, n_features), target Y has shape (n_samples, 1)
Use tf.keras.layers.Dense(1) for the simplest Keras linear regression
Use tf.GradientTape for manual gradient descent (educational purposes)
Use the normal equation W = (X^T X)^(-1) X^T Y for small datasets
Always normalize features before training for faster convergence
Ensure both X and Y are float32 and properly shaped