The Frameworks

The Frameworks

Share this post

The Frameworks
The Frameworks
Building Your First Supervised Learning Model in TensorFlow
User's avatar
Discover more from The Frameworks
Stories and insights on engineering, data, privacy, and the human side of tech.
Already have an account? Sign in

Building Your First Supervised Learning Model in TensorFlow

Machine Learning 101 - Lesson 5

Johnny Unar's avatar
Johnny Unar
Mar 08, 2025

Share this post

The Frameworks
The Frameworks
Building Your First Supervised Learning Model in TensorFlow
Share

Previous lesson:

Supervised vs. Unsupervised Learning: What’s the Difference?

Supervised vs. Unsupervised Learning: What’s the Difference?

Johnny Unar
·
Mar 4
Read full story

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?

  1. Data Preparation → We load and preprocess the Titanic dataset by encoding categorical features and scaling numerical data.

  2. Model Definition → We build a simple neural network with two hidden layers.

  3. Compilation → We use Adam optimizer and binary cross-entropy loss.

    1. Adam optimizer is a popular optimization algorithm that adjusts learning rates during training, making it efficient for deep learning.

    2. Binary cross-entropy loss measures how well the predicted probability aligns with actual labels in binary classification.

  4. Training → The model learns patterns in the training data over 20 epochs.

  5. Evaluation → We test how well the model generalizes to unseen passengers.

  6. 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! 🚀


Subscribe to The Frameworks

By Johnny Unar · Launched 7 months ago
Stories and insights on engineering, data, privacy, and the human side of tech.

Share this post

The Frameworks
The Frameworks
Building Your First Supervised Learning Model in TensorFlow
Share

Discussion about this post

User's avatar
We Are Standing On The Edge (Let’s not f*ck this up)
This is a letter to all the folks who need to hear this, to all the people who feel lost and overpowered by the algorithms thinking for them, for us.
Dec 15, 2024 â€¢ 
Johnny Unar
3

Share this post

The Frameworks
The Frameworks
We Are Standing On The Edge (Let’s not f*ck this up)
Český Tarzan: Přirozený pohyb, biohacking a síla mysli
Vítej u mého historicky prvního a s trochou štěstí snad ne posledního rozhovoru na této platformě.
Feb 22 â€¢ 
Johnny Unar
1

Share this post

48:00
Why Coding Is Even More Relevant For You In The AI Era: A Story
I have been in the software engineering space for quite some time now and this letter is for anyone who thinks there is no point even trying.
Dec 29, 2024 â€¢ 
Johnny Unar
1

Share this post

The Frameworks
The Frameworks
Why Coding Is Even More Relevant For You In The AI Era: A Story

Ready for more?

© 2025 Jan Unar
Privacy ∙ Terms ∙ Collection notice
Start writingGet the app
Substack is the home for great culture

Share

Create your profile

User's avatar

Only paid subscribers can comment on this post

Already a paid subscriber? Sign in

Check your email

For your security, we need to re-authenticate you.

Click the link we sent to , or click here to sign in.