Linear Regression¶

In this notebook, we will train a linear regression model using PyTorch to predict the median house value in California. We will start from a simple model, then improve its accuracy by cleaning our data and engineering new features.

To learn the theoretical concepts involved in Linear Regression, read this doc.

In [1]:

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.datasets import fetch_california_housing
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset

import matplotlib.pyplot as plt import numpy as np import pandas as pd from sklearn.datasets import fetch_california_housing from sklearn.metrics import r2_score from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torch.utils.data import TensorDataset

What Are We Importing?¶

We need three groups of libraries:

• Data handling: pandas and numpy for loading and manipulating the dataset.
• Scikit-learn utilities: fetch_california_housing to download the data, train_test_split to split it, StandardScaler to normalize features, and r2_score to evaluate our model.
• PyTorch: torch for tensors, nn for the linear layer and loss function, optim for the Adam optimizer, and DataLoader/TensorDataset for feeding the data to the model in small batches.

In [2]:

housing = fetch_california_housing()

housing = fetch_california_housing()

The California Housing Dataset¶

fetch_california_housing() returns a Bunch object containing:

• housing["data"]: the features — 8 numerical attributes about each neighborhood, such as median income (MedInc), house age, and geographic coordinates.
• housing["target"]: the label — the median house value (MedHouseVal) we want to predict.
• housing["feature_names"] and housing["target_names"]: the names of the columns.

We combine everything into a single pandas DataFrame so we can inspect and clean it easily. A DataFrame is a tabular structure — rows are examples, columns are features and the label.

In [3]:

df = pd.DataFrame(data=housing["data"], columns=housing["feature_names"])
df[housing["target_names"][0]] = housing["target"]

df = pd.DataFrame(data=housing["data"], columns=housing["feature_names"]) df[housing["target_names"][0]] = housing["target"]

Why Clean the Data First?¶

In [4]:

df.describe()

df.describe()

Out[4]:

	MedInc	HouseAge	AveRooms	AveBedrms	Population	AveOccup	Latitude	Longitude	MedHouseVal
count	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000
mean	3.870671	28.639486	5.429000	1.096675	1425.476744	3.070655	35.631861	-119.569704	2.068558
std	1.899822	12.585558	2.474173	0.473911	1132.462122	10.386050	2.135952	2.003532	1.153956
min	0.499900	1.000000	0.846154	0.333333	3.000000	0.692308	32.540000	-124.350000	0.149990
25%	2.563400	18.000000	4.440716	1.006079	787.000000	2.429741	33.930000	-121.800000	1.196000
50%	3.534800	29.000000	5.229129	1.048780	1166.000000	2.818116	34.260000	-118.490000	1.797000
75%	4.743250	37.000000	6.052381	1.099526	1725.000000	3.282261	37.710000	-118.010000	2.647250
max	15.000100	52.000000	141.909091	34.066667	35682.000000	1243.333333	41.950000	-114.310000	5.000010

If we look closely at df.describe(), the max values tell an important story:

Feature	75th percentile	Max
AveRooms	6.05	141.9
AveOccup	3.28	1243.3

The vast majority of neighborhoods have a handful of rooms and occupants, yet a few extreme values reach into the hundreds or thousands. These are outliers — extreme values that sit far outside the normal range. We will handle them before training.

Cleaning the Data: Clipping Outliers¶

A linear model is trained using Mean Squared Error, which squares the error for every example. A single neighborhood with 141 rooms would produce an enormous squared error, pulling the model's weights in a weird direction just to accommodate that one row.

The fix: we clip each column at its 99th percentile. quantile(0.99) returns the value below which 99% of the data falls, and clip(upper=...) caps anything above that value. This tames the extreme tail without removing any rows.

To see why this helps, imagine every neighborhood sorted by room count, from smallest to largest. The 99th percentile is the line that separates the normal range from the extreme tail:

Sorted houses:  [ 1.1, 2.5, 4.0, ... , 13.1, 13.2 ]   [ 45.2, 89.0, 141.9 ]
                \_________________________________/    \_________________/
                   99% of houses (normal range)         Top 1% (outliers)
                                                    ^
                                99th percentile (quantile 0.99)

The problem: without clipping, those few outliers stretch the number line so far that all the normal houses get squished together:

               Normal houses                          Outliers
             [||||||||||||||||||||||||||]                    |         |
+------------+---------------------------+-------------------+---------+-----------> rooms
0           10                          30                  60       100         150

After clipping at the 99th percentile, the outliers are pulled back to the cap, so the model can actually tell normal houses apart:

                Normal houses
              [||||||||||||||||||||||||||]   <-- outliers are now capped here
+------------+---------------------------+-------------------+---------+-----------> rooms
0           10                   13.2 (cap)                  60       100         150

We also engineer two ratio features:

• RoomsPerHousehold = AveRooms / AveOccup — how spacious a home is per person.
• BedroomsPerRoom = AveBedrms / AveRooms — how much of the home is bedrooms.

These give the model information it cannot derive from the raw columns alone. Finally, we split the data into train and test sets using a fixed random_state so our results are reproducible.

In [5]:

df_clean = df.copy()
for c in ["AveRooms", "AveBedrms", "Population", "AveOccup"]:
  upper = df_clean[c].quantile(0.99)
  df_clean[c] = df_clean[c].clip(upper=upper)

df_clean["RoomsPerHousehold"] = df_clean["AveRooms"] / df_clean["AveOccup"]
df_clean["BedroomsPerRoom"] = df_clean["AveBedrms"] / df_clean["AveRooms"]

feature_names_updated = housing["feature_names"] + [
    "RoomsPerHousehold", "BedroomsPerRoom"
]

X_numpy = df_clean[feature_names_updated].to_numpy()
y_numpy = df_clean[housing["target_names"][0]].to_numpy()

df_clean = df.copy() for c in ["AveRooms", "AveBedrms", "Population", "AveOccup"]: upper = df_clean[c].quantile(0.99) df_clean[c] = df_clean[c].clip(upper=upper) df_clean["RoomsPerHousehold"] = df_clean["AveRooms"] / df_clean["AveOccup"] df_clean["BedroomsPerRoom"] = df_clean["AveBedrms"] / df_clean["AveRooms"] feature_names_updated = housing["feature_names"] + [ "RoomsPerHousehold", "BedroomsPerRoom" ] X_numpy = df_clean[feature_names_updated].to_numpy() y_numpy = df_clean[housing["target_names"][0]].to_numpy()

In [6]:

df_clean.describe()

df_clean.describe()

Out[6]:

	MedInc	HouseAge	AveRooms	AveBedrms	Population	AveOccup	Latitude	Longitude	MedHouseVal	RoomsPerHousehold	BedroomsPerRoom
count	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000	20640.000000
mean	3.870671	28.639486	5.330588	1.076287	1403.613896	2.915167	35.631861	-119.569704	2.068558	1.936423	0.213026
std	1.899822	12.585558	1.330038	0.160058	973.476399	0.734751	2.135952	2.003532	1.153956	0.644395	0.057784
min	0.499900	1.000000	0.846154	0.333333	3.000000	0.692308	32.540000	-124.350000	0.149990	0.311943	0.100000
25%	2.563400	18.000000	4.440716	1.006079	787.000000	2.429741	33.930000	-121.800000	1.196000	1.523082	0.175662
50%	3.534800	29.000000	5.229129	1.048780	1166.000000	2.818116	34.260000	-118.490000	1.797000	1.937936	0.203690
75%	4.743250	37.000000	6.052381	1.099526	1725.000000	3.282261	37.710000	-118.010000	2.647250	2.296090	0.239466
max	15.000100	52.000000	10.357033	2.127541	5805.830000	5.394812	41.950000	-114.310000	5.000010	14.960159	1.000000

The Train / Test Split¶

train_test_split shuffles the data and reserves 33.3% of it as the test set. The model will only ever see the training portion during training. We keep the test set completely hidden until the very end, so the final score tells us how well the model generalizes to data it has never seen.

In [7]:

X_train, X_test, y_train, y_test = train_test_split(X_numpy,
                                                    y_numpy,
                                                    test_size=0.333,
                                                    random_state=42)

X_train, X_test, y_train, y_test = train_test_split(X_numpy, y_numpy, test_size=0.333, random_state=42)

Scaling the Features¶

Our features are on very different scales — MedInc is in the single digits, while Population is in the thousands. Without scaling, a feature with larger numbers would dominate the loss purely because of its magnitude. This is called feature dominance.

StandardScaler subtracts the mean and divides by the standard deviation of each column, so every feature is centered around 0 with unit variance. We call fit_transform on the training data only, then transform the test data using the same statistics — never fitting on the test set.

In [8]:

scaler = StandardScaler()
X_train, X_test = scaler.fit_transform(X_train), scaler.transform(X_test)

scaler = StandardScaler() X_train, X_test = scaler.fit_transform(X_train), scaler.transform(X_test)

From NumPy to Tensors¶

PyTorch cannot work directly with NumPy arrays, so we convert them into tensors — PyTorch's n-dimensional arrays. We cast to float32 because that is the default precision PyTorch expects, and we reshape y into a column vector of shape (n_samples, 1) to match the model's output shape. y.view(-1, 1) is shorthand for y.view(y.shape[0], 1).

In [9]:

X = torch.from_numpy(X_train.astype(np.float32))
y = torch.from_numpy(y_train.astype(np.float32))
y = y.view(-1, 1)  # equivalent to `y.view(y.shape[0], 1)`

X = torch.from_numpy(X_train.astype(np.float32)) y = torch.from_numpy(y_train.astype(np.float32)) y = y.view(-1, 1) # equivalent to `y.view(y.shape[0], 1)`

Batching the Data: The DataLoader¶

Computing the gradient over all 13,766 training examples at once is slow and noisy. Instead, we use mini-batches. TensorDataset pairs each X sample with its y label, and DataLoader iterates over shuffled chunks of batch_size=64. This gives the optimizer many small weight updates per epoch instead of one giant one, which converges faster and more stably.

We keep the test tensors separate — they are not batched because we only need one forward pass through them.

In [10]:

train_dataset = TensorDataset(X, y)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

train_dataset = TensorDataset(X, y) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

In [11]:

X_test = torch.from_numpy(X_test.astype(np.float32))
y_test = torch.from_numpy(y_test.astype(np.float32))
y_test = y_test.view(-1, 1)

X_test = torch.from_numpy(X_test.astype(np.float32)) y_test = torch.from_numpy(y_test.astype(np.float32)) y_test = y_test.view(-1, 1)

In [12]:

n_samples, n_features = X.shape
n_samples, n_features

n_samples, n_features = X.shape n_samples, n_features

Out[12]:

(13766, 10)

Building the Model¶

nn.Linear(n_features, 1) creates a single layer that computes $y = XW + b$, where $W$ is a weight matrix of shape (n_features, 1) and $b$ is a scalar bias. Because we engineered two extra features, n_features is 10 (not 8) — notice the shape (13766, 10) from the previous cell. The model therefore has 10 weights plus a bias, for 11 learnable parameters in total.

In [13]:

model = nn.Linear(n_features, 1)

model = nn.Linear(n_features, 1)

Loss Function and Optimizer¶

• nn.MSELoss() — the Mean Squared Error loss. It measures the average squared difference between predictions and true values. Perfect predictions give a loss of 0; larger errors are punished quadratically.
• optim.Adam(...) — an optimizer that updates the weights using gradient descent. Adam adapts the learning rate for each parameter individually, which usually converges faster than plain SGD. The lr=0.01 controls the step size of each update.

In [14]:

criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.01)

The Training Loop¶

Each epoch does the following:

1. Forward pass — predict house values for every batch with model(X_batch).
2. Compute loss — compare predictions to the true labels using MSE.
3. Zero the gradients — optimizer.zero_grad() clears the accumulated gradients from the previous step, otherwise they would add up across batches.
4. Backward pass — loss.backward() computes the gradient of the loss with respect to every parameter.
5. Update — optimizer.step() nudges the weights in the direction that reduces the loss.

We track the average loss per epoch so we can watch the model improve. Notice that the loss drops sharply in the first few epochs and then plateaus around 0.44 — the model has converged.

In [15]:

n_epochs = 30
for epoch in range(n_epochs):
  model.train()
  total_loss = 0.0

  for X_batch, y_batch in train_loader:
    y_pred = model(X_batch)
    loss = criterion(y_pred, y_batch)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    total_loss += loss.item() * len(X_batch)

  avg_loss = total_loss / len(X)
  if (epoch + 1) % 5 == 0:
    print(f"Epoch {epoch+1:2d} | Train Loss: {avg_loss:.4f}")

n_epochs = 30 for epoch in range(n_epochs): model.train() total_loss = 0.0 for X_batch, y_batch in train_loader: y_pred = model(X_batch) loss = criterion(y_pred, y_batch) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() * len(X_batch) avg_loss = total_loss / len(X) if (epoch + 1) % 5 == 0: print(f"Epoch {epoch+1:2d} | Train Loss: {avg_loss:.4f}")

Epoch  5 | Train Loss: 0.4368
Epoch 10 | Train Loss: 0.4415
Epoch 15 | Train Loss: 0.4376
Epoch 20 | Train Loss: 0.4381
Epoch 25 | Train Loss: 0.4386
Epoch 30 | Train Loss: 0.4381

Evaluating the Model¶

After training, we run the model on the held-out test set and measure performance with the coefficient of determination $R^2$.

$R^2$ tells us what fraction of the variance in house prices our model explains, on a scale from negative values up to 1.0 (a perfect fit). We detach the predictions from the computation graph (.detach()) and convert them back to NumPy before scoring.

In [16]:

y_pred = model(X_test).detach().numpy()
y_pred = y_pred.flatten()

y_pred = model(X_test).detach().numpy() y_pred = y_pred.flatten()

In [17]:

r2_score(y_test, y_pred)

r2_score(y_test, y_pred)

Out[17]:

0.6684893369674683

What Did We Achieve?¶

Our test $R^2$ score is around 0.67, up from roughly 0.59 on the raw features. The two changes that mattered:

1. Clipping outliers at the 99th percentile stopped a handful of extreme neighborhoods from distorting the loss.
2. Engineering ratio features gave the model relationships it could not discover from the raw columns alone.

A purely linear model has its limits though. To push higher, we could add polynomial features to capture non-linear interactions — a natural next experiment.

Copy page as MarkdownCopy to clipboard for pasting into any AI tool View as MarkdownOpen this page as clean Markdown in a new tab

Open in ChatGPTChat about this page with ChatGPT Open in ClaudeChat about this page with Claude

llms.txtFull documentation index for LLMs

Use with AI