< Home
Week 4: Assignment ~ Training a Model¶
In this assignment I have trained a linear regression model to predict house price (in $1000s) from size (sqft), leveraging the clear linear relationship in the data.
It is a Supervised Learning: It has labeled data — each input (house size) has a corresponding output (price). The goal is to learn a function f such that:
price ≈ f(size)Training a Model¶
In [1]:
import jax.numpy as jnp
import matplotlib.pyplot as plt
# Simple house data: [size_in_sqft, price_in_thousands]
house_data = jnp.array([
[1000, 200], # 1000 sqft → $200,000
[1500, 300], # 1500 sqft → $300,000
[2000, 400], # 2000 sqft → $400,000
[2500, 500], # 2500 sqft → $500,000
[3000, 600], # 3000 sqft → $600,000
[1200, 240], # 1200 sqft → $240,000
[1800, 360], # 1800 sqft → $360,000
], dtype=jnp.float32)
# Split into inputs (size) and outputs (price)
X = house_data[:, 0:1] # Size column
y = house_data[:, 1:2] # Price column
print("House Data:")
print("Size (sqft) -> Price ($1000s)")
for i in range(len(X)):
print(f"{X[i][0]:.0f} -> {y[i][0]:.0f}")
# Let's plot it to see the pattern
plt.scatter(X, y)
plt.xlabel('House Size (sqft)')
plt.ylabel('Price ($1000s)')
plt.title('House Prices vs Size')
plt.show()
House Data: Size (sqft) -> Price ($1000s) 1000 -> 200 1500 -> 300 2000 -> 400 2500 -> 500 3000 -> 600 1200 -> 240 1800 -> 360
The second one is same model like above but with small modification where after training, it can predict the house price at real time as per the user input.
In [ ]:
import jax.numpy as jnp
import jax
from jax import grad
import matplotlib.pyplot as plt
# 🔹 House data: [size_in_sqft, price_in_thousands]
house_data = jnp.array([
[1000, 200],
[1500, 300],
[2000, 400],
[2500, 500],
[3000, 600],
[1200, 240],
[1800, 360],
], dtype=jnp.float32)
X = house_data[:, 0:1] # (n, 1)
y = house_data[:, 1:2] # (n, 1)
print("House Data:")
print("Size (sqft) -> Price ($1000s)")
for i in range(len(X)):
print(f"{X[i][0]:.0f} -> {y[i][0]:.0f}")
# 🔹 Normalize features (optional but helps convergence)
X_mean = X.mean()
X_std = X.std()
X_norm = (X - X_mean) / X_std # Z-score normalization
# 🔹 Initialize model: y_pred = w * x + b
key = jax.random.PRNGKey(0)
w = jax.random.normal(key, (1, 1)) * 0.1 # small init
b = jnp.array([0.0])
# 🔹 Loss function: Mean Squared Error
def loss_fn(w, b, X, y):
y_pred = X @ w + b # (n,1) = (n,1)@(1,1) + (1,)
return jnp.mean((y_pred - y) ** 2)
# 🔹 Compute gradients
grad_w = grad(loss_fn, argnums=0)
grad_b = grad(loss_fn, argnums=1)
# 🔹 Training loop
learning_rate = 0.01
epochs = 500
for epoch in range(epochs):
loss_val = loss_fn(w, b, X_norm, y)
dw = grad_w(w, b, X_norm, y)
db = grad_b(w, b, X_norm, y)
w -= learning_rate * dw
b -= learning_rate * db
if epoch % 100 == 0:
print(f"Epoch {epoch:3d} | Loss: {loss_val:.4f} | w: {w[0,0]:.4f}, b: {b[0]:.4f}")
print(f"\n✅ Training complete.")
print(f"Learned model (on normalized input): price = {w[0,0]:.4f} * (size_norm) + {b[0]:.4f}")
print(f"Denormalized slope ≈ {w[0,0] / X_std:.4f} → expected ~0.2")
print(f"True relationship: price = 0.2 * size")
# 🔹 Function to predict (handles normalization internally)
def predict_price(size_sqft):
size_sqft = jnp.array(size_sqft, dtype=jnp.float32).reshape(-1, 1)
size_norm = (size_sqft - X_mean) / X_std
pred = size_norm @ w + b
return pred.squeeze() # scalar or array in $1000s
# 🔹 Plot original data + learned line
plt.figure(figsize=(8, 5))
plt.scatter(X, y, color='blue', label='Data')
# Generate smooth line for visualization
x_plot = jnp.linspace(X.min(), X.max(), 100).reshape(-1, 1)
y_plot = predict_price(x_plot)
plt.plot(x_plot, y_plot, color='red', label='Learned model', linewidth=2)
plt.xlabel('House Size (sqft)')
plt.ylabel('Price ($1000s)')
plt.title('House Price Prediction (Linear Regression)')
plt.legend()
plt.grid(True, linestyle='--', alpha=0.6)
plt.show()
# 🔹 Real-time prediction loop
print("\n" + "="*50)
print("🏡 Real-Time House Price Prediction")
print("Enter house size in square feet (or 'q' to quit):")
print("="*50)
while True:
try:
user_input = input("\nEnter size (sqft): ").strip()
if user_input.lower() in ['q', 'quit', 'exit']:
print("👋 Exiting. Thanks for using the predictor!")
break
size = float(user_input)
if size <= 0:
print("⚠️ Please enter a positive number.")
continue
price_pred = predict_price(size)
print(f"✅ Predicted price: ${price_pred:,.0f},000")
# e.g., 2200 → $440,000
except ValueError:
print("⚠️ Invalid input. Please enter a number or 'q' to quit.")
except KeyboardInterrupt:
print("\n\n👋 Interrupted. Goodbye!")
break
House Data: Size (sqft) -> Price ($1000s) 1000 -> 200 1500 -> 300 2000 -> 400 2500 -> 500 3000 -> 600 1200 -> 240 1800 -> 360 Epoch 0 | Loss: 155271.5625 | w: 2.7938, b: 7.4286 Epoch 100 | Loss: 2730.9084 | w: 114.6381, b: 323.1551 Epoch 200 | Loss: 48.0295 | w: 129.4709, b: 365.0267 Epoch 300 | Loss: 0.8448 | w: 131.4380, b: 370.5795 Epoch 400 | Loss: 0.0148 | w: 131.6989, b: 371.3160 ✅ Training complete. Learned model (on normalized input): price = 131.7334 * (size_norm) + 371.4133 Denormalized slope ≈ 0.2000 → expected ~0.2 True relationship: price = 0.2 * size
================================================== 🏡 Real-Time House Price Prediction Enter house size in square feet (or 'q' to quit): ==================================================
✅ Predicted price: $1,000,000
✅ Predicted price: $400,000
✅ Predicted price: $20,000
✅ Predicted price: $200,000
In [15]:
import jax
import jax.numpy as jnp
from jax import random, grad, jit
# Initialize random key
key = random.PRNGKey(42)
# We're using the house data from above
# X = house sizes, y = house prices
# Parameter initialization - MODIFIED for 1 input, 1 output
def init_params(key):
k1, k2 = random.split(key)
W1 = 0.1 * random.normal(k1, (1, 4)) # 1 input → 4 neurons
b1 = jnp.zeros(4)
W2 = 0.1 * random.normal(k2, (4, 1)) # 4 neurons → 1 output
b2 = jnp.zeros(1)
return (W1, b1, W2, b2)
# Forward pass - SAME as before but we'll keep it linear
def forward(params, in0):
W1, b1, W2, b2 = params
h1 = jnp.tanh(in0 @ W1 + b1) # Hidden layer with tanh
o2 = h1 @ W2 + b2 # Output layer - LINEAR (no sigmoid) for prices
return o2
# Loss function - SAME as before
def loss(params):
y_pred = forward(params, X)
return jnp.mean((y_pred - y) ** 2)
# Update weights - SAME as before
@jit
def update(params, rate=0.001): # Smaller learning rate for prices
grads = grad(loss)(params)
return jax.tree.map(lambda p, g: p - rate * g, params, grads)
# Initialize and train
params = init_params(key)
print("Training the house price predictor...")
for step in range(10000):
params = update(params, rate=0.001)
if step % 400 == 0:
current_loss = loss(params)
print(f"Step {step:4d}, Loss: {current_loss:.2f}")
# Test the model
print("\n=== PREDICTIONS ===")
print("Size -> Actual Price | Predicted Price")
predictions = forward(params, X)
for i in range(len(X)):
actual = y[i][0]
predicted = predictions[i][0]
error = abs(actual - predicted)
print(f"{X[i][0]:.0f} sqft -> ${actual:.0f}K | ${predicted:.0f}K (error: ${error:.1f}K)")
# Predict a NEW house size
new_size = jnp.array([[1750.0]]) # A house we haven't seen
predicted_price = forward(params, new_size)
print(f"\nPrediction for {new_size[0][0]:.0f} sqft house: ${predicted_price[0][0]:.0f}K")
Training the house price predictor... Step 0, Loss: 152695.44 Step 400, Loss: 17398.71 Step 800, Loss: 17355.12 Step 1200, Loss: 17355.10 Step 1600, Loss: 17355.10 Step 2000, Loss: 17355.10 Step 2400, Loss: 17355.10 Step 2800, Loss: 17355.10 Step 3200, Loss: 17355.10 Step 3600, Loss: 17355.10 Step 4000, Loss: 17355.10 Step 4400, Loss: 17355.10 Step 4800, Loss: 17355.10 Step 5200, Loss: 17355.10 Step 5600, Loss: 17355.10 Step 6000, Loss: 17355.10 Step 6400, Loss: 17355.10 Step 6800, Loss: 17355.10 Step 7200, Loss: 17355.10 Step 7600, Loss: 17355.10 Step 8000, Loss: 17355.10 Step 8400, Loss: 17355.10 Step 8800, Loss: 17355.10 Step 9200, Loss: 17355.10 Step 9600, Loss: 17355.10 === PREDICTIONS === Size -> Actual Price | Predicted Price 1000 sqft -> $200K | $371K (error: $171.4K) 1500 sqft -> $300K | $371K (error: $71.4K) 2000 sqft -> $400K | $371K (error: $28.6K) 2500 sqft -> $500K | $371K (error: $128.6K) 3000 sqft -> $600K | $371K (error: $228.6K) 1200 sqft -> $240K | $371K (error: $131.4K) 1800 sqft -> $360K | $371K (error: $11.4K) Prediction for 1750 sqft house: $371K
In [17]:
# Test on various house sizes
test_sizes = jnp.array([[800], [1600], [2200], [2800], [3500]])
test_predictions = forward(params, test_sizes)
print("\n=== PRICE PREDICTIONS FOR DIFFERENT SIZES ===")
for i in range(len(test_sizes)):
size = test_sizes[i][0]
price = test_predictions[i][0]
print(f"{size:.0f} sqft → ${price:.0f}K"
# Plot the results
plt.figure(figsize=(10, 5))
# Plot training data
plt.scatter(X, y, color='blue', label='Actual Prices', s=50)
# Plot predictions
all_sizes = jnp.concatenate([X, test_sizes])
all_prices = forward(params, all_sizes)
plt.plot(all_sizes, all_prices, color='red', label='Model Predictions', linewidth=2)
plt.xlabel('House Size (sqft)')
plt.ylabel('Price ($1000s)')
plt.title('House Price Predictor')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()
=== PRICE PREDICTIONS FOR DIFFERENT SIZES === 800 sqft → $371K 1600 sqft → $371K 2200 sqft → $371K 2800 sqft → $371K 3500 sqft → $371K
In [ ]: