
Previous lesson:
Now that we understand supervised learning, it's time to build our first supervised learning model using TensorFlow. In this lesson, we’ll train a model to perform classification, one of the most common tasks in machine learning.
By the end of this lesson, you’ll:
Load and preprocess a dataset.
Build a neural network in TensorFlow.
Train the model on labeled data.
Evaluate its performance and make predictions.
Let’s get started! 🚀
Dataset: Predicting Titanic Survival
We’ll use the Titanic dataset, which contains information about passengers on the Titanic, including their age, gender, ticket class, and whether they survived. Our goal is to build a classification model that predicts who survived based on these features.
Step 1: Import Dependencies and Load the Dataset
First, let’s import the necessary libraries and load the dataset.
import tensorflow as tf
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
# Load the Titanic dataset
data = pd.read_csv("https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv")
# Select relevant features and target
features = ["Pclass", "Sex", "Age", "SibSp", "Parch", "Fare"]
target = "Survived"
data = data[features + [target]].dropna() # Remove missing values
Step 2: Preprocess the Data
We need to convert categorical data into numerical values and scale numerical features.
# Encode categorical feature 'Sex'
encoder = LabelEncoder()
data["Sex"] = encoder.fit_transform(data["Sex"])
# Split into features and labels
X = data[features]
y = data[target]
# Standardize numerical features
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
# Split into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
Step 3: Build the Neural Network
Now, we define our model. We'll use a simple feedforward neural network with:
An input layer that takes passenger features.
A hidden layer with 16 neurons and ReLU activation.
ReLU (Rectified Linear Unit) is a commonly used activation function that outputs the input value if it’s positive and zero otherwise, helping the network learn complex patterns efficiently.
An output layer with 1 neuron and sigmoid activation for binary classification.
Sigmoid activation outputs a probability between 0 and 1, making it ideal for binary classification problems like survival prediction.
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(X_train.shape[1],)),
tf.keras.layers.Dense(16, activation='relu'),
tf.keras.layers.Dense(8, activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid') # Output layer for binary classification
])
# Compile the model
model.compile(optimizer='adam',
loss='binary_crossentropy',
metrics=['accuracy'])
Step 4: Train the Model
Let’s train our model using the training dataset. We’ll train it for 50 epochs.
An epoch is one complete pass through the entire training dataset. Training for multiple epochs helps the model gradually learn and improve its predictions.
# Train the model
model.fit(X_train, y_train, epochs=50, validation_data=(X_test, y_test))
During training, TensorFlow will display the loss and accuracy for each epoch. A lower loss and higher accuracy mean the model is improving.
Step 5: Evaluate and Make Predictions
After training, we evaluate the model’s accuracy on unseen test data:
# Evaluate the model on test data
test_loss, test_acc = model.evaluate(X_test, y_test)
print(f"Test accuracy: {test_acc:.4f}")
Now, let's use the model to make a prediction for a new passenger:
# Example passenger: (Pclass=3, Sex=female, Age=22, SibSp=1, Parch=0, Fare=7.25)
new_passenger = np.array([[3, 1, 22, 1, 0, 7.25]])
new_passenger_scaled = scaler.transform(new_passenger)
# Predict survival
prediction = model.predict(new_passenger_scaled)
print(f"Survival Probability: {prediction[0][0]:.4f}")
If the predicted probability is above 0.5, the model predicts that the passenger survived, otherwise, it predicts they did not survive.
Why 0.5?
The 0.5 threshold is a standard choice for binary classification problems where the outputs are probabilities:
A probability closer to 1 means the model is very confident the passenger survived.
A probability closer to 0 means the model is very confident the passenger did not survive.
If the probability is exactly 0.5, the model is completely uncertain and could go either way.
What’s Happening Here?
Data Preparation → We load and preprocess the Titanic dataset by encoding categorical features and scaling numerical data.
Model Definition → We build a simple neural network with two hidden layers.
Compilation → We use Adam optimizer and binary cross-entropy loss.
Adam optimizer is a popular optimization algorithm that adjusts learning rates during training, making it efficient for deep learning.
Binary cross-entropy loss measures how well the predicted probability aligns with actual labels in binary classification.
Training → The model learns patterns in the training data over 20 epochs.
Evaluation → We test how well the model generalizes to unseen passengers.
Prediction → We feed in a new passenger’s details to predict survival probability.
Full Code
import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
def load_and_preprocess_data(url: str) -> tuple[np.ndarray, pd.Series, StandardScaler]:
"""
Load the Titanic dataset from a URL, preprocess by encoding categorical features,
and standardize numerical features.
Args:
url (str): URL of the dataset.
Returns:
tuple: A tuple containing the scaled features array, target series, and the fitted scaler.
"""
data = pd.read_csv(url)
features = ["Pclass", "Sex", "Age", "SibSp", "Parch", "Fare"]
target = "Survived"
data = data[features + [target]].dropna()
# Encode categorical feature 'Sex'
encoder = LabelEncoder()
data["Sex"] = encoder.fit_transform(data["Sex"])
X = data[features]
y = data[target]
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
return X_scaled, y, scaler
def build_model(input_dim: int) -> tf.keras.Model:
"""
Build and compile a Keras Sequential model for binary classification.
Args:
input_dim (int): Number of features in the input.
Returns:
tf.keras.Model: The compiled Keras model.
"""
model = tf.keras.Sequential(
[
tf.keras.layers.Input(shape=(input_dim,)),
tf.keras.layers.Dense(16, activation="relu"),
tf.keras.layers.Dense(8, activation="relu"),
tf.keras.layers.Dense(1, activation="sigmoid"),
]
)
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
return model
def main() -> None:
"""
Load data, train and evaluate the model, and make a prediction.
"""
data_url = (
"https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv"
)
X_scaled, y, scaler = load_and_preprocess_data(data_url)
# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(
X_scaled, y, test_size=0.2, random_state=42
)
model = build_model(X_train.shape[1])
model.fit(X_train, y_train, epochs=50, validation_data=(X_test, y_test))
test_loss, test_acc = model.evaluate(X_test, y_test)
print(f"Test accuracy: {test_acc:.4f}")
# Define the features used during training
features = ["Pclass", "Sex", "Age", "SibSp", "Parch", "Fare"]
# Create a DataFrame for the new passenger to include valid feature names
new_passenger = np.array([[3, 1, 22, 1, 0, 7.25]])
new_passenger_df = pd.DataFrame(new_passenger, columns=features)
new_passenger_scaled = scaler.transform(new_passenger_df)
prediction = model.predict(new_passenger_scaled)
print(f"Survival Probability: {prediction[0][0]:.4f}")
if __name__ == "__main__":
main()
Next Steps
Congratulations! 🎉 You’ve built a supervised learning model to predict survival on the Titanic using TensorFlow. In the next lesson, we’ll explore improving model performance by tuning hyperparameters and adding techniques like dropout and batch normalization. Stay tuned! 🚀