Index Ai Academy
Project tutorials · 7 steps

Train a neural network classifier in PyTorch

Prepare data without leaks, run the checks that catch bugs, train with early stopping, and compare honestly with a simple baseline.

What you'll build

Training a network is easy. Knowing whether to trust it is the hard part. This tutorial trains a classifier on tabular churn data with the checks careful engineers run first, then scores it against logistic regression. The simple model wins, and that's a useful result to see for yourself.

What you'll learn

  • Prepare tabular data for a neural network without leaking test information
  • Run the sanity checks before trusting a training run
  • Train a regularized network with early stopping on a validation metric
  • Compare it honestly with a classical baseline, and know when deep learning pays off

What you need

Python, pandas, and what a neural network and a loss are. In the lesson it runs in your browser, on an in-browser subset of PyTorch. On your own machine, install PyTorch, pandas and scikit-learn, and change the data address to https://www.indexaiacademy.com/mock/v1/customers. Real PyTorch gives slightly different numbers.

Step 1: Prepare the data without leaking

prepare learns medians, means and scales from the training set only, then applies them everywhere. Using the whole dataset would quietly leak information about the test set into training. email_missing records where a value was absent, which can itself be a signal.

import math
import numpy as np
import pandas as pd
import requests
import torch
import torch.nn as nn
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

df = pd.DataFrame(requests.get("https://api.gradient.test/v1/customers").json()["data"])
numeric = ["tenure_months", "orders_last_year", "avg_order_value", "days_since_last_order",
           "discount_share", "email_open_rate", "support_tickets", "trail_club"]
categorical = ["region", "channel", "top_category"]
df["email_missing"] = df["email_open_rate"].isna().astype(float)
dev, test = train_test_split(df, test_size=0.2, random_state=7, stratify=df["churned"])
train, val = train_test_split(dev, test_size=0.25, random_state=7, stratify=dev["churned"])

def prepare(fit, *frames):
    """Learn medians, means, scales and categories from `fit` only, then apply them to every frame."""
    med = fit[numeric].median()
    mu, sd = fit[numeric].fillna(med).mean(), fit[numeric].fillna(med).std()
    cats = pd.get_dummies(fit[categorical].replace("", "unknown")).columns
    out = []
    for f in frames:
        x = (f[numeric].fillna(med) - mu) / sd
        d = pd.get_dummies(f[categorical].replace("", "unknown")).reindex(columns=cats, fill_value=0)
        out.append(torch.tensor(np.c_[x.to_numpy(float), f[["email_missing"]].to_numpy(float), d.to_numpy(float)], dtype=torch.float32))
    return out

X_train, X_val, X_test = prepare(train, train, val, test)
y_train, y_val, y_test = (torch.tensor(f["churned"].to_numpy()) for f in (train, val, test))
print("shapes:", tuple(X_train.shape), tuple(X_val.shape), tuple(X_test.shape), f"  churn rate {y_train.float().mean().item():.1%}")
What it prints
shapes: (360, 20) (120, 20) (120, 20)   churn rate 19.4%

Step 2: Check the starting loss

An untrained network should guess about evenly, and with two classes that gives a loss near ln 2. A very different number means a bug in the labels, the loss or the setup, and it takes one second to find out.

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(X_train.shape[1], 32), nn.ReLU(), nn.Linear(32, 2))
with torch.no_grad():
    start = nn.functional.cross_entropy(model(X_train), y_train).item()
print(f"starting loss {start:.3f}   (two classes: ln 2 = {math.log(2):.3f})")
What it prints
starting loss 0.704   (two classes: ln 2 = 0.693)

Step 3: Prove it can memorize a tiny batch

If a network can't memorize 16 examples, it has a bug, and more data won't help. The loss should fall almost to zero, and the check takes seconds.

opt = torch.optim.Adam(model.parameters(), lr=0.01)       # can it memorize 16 customers?
for _ in range(300):
    opt.zero_grad()
    loss = nn.functional.cross_entropy(model(X_train[:16]), y_train[:16])
    loss.backward()
    opt.step()
print(f"loss after overfitting 16 examples: {loss.item():.4f}")
What it prints
loss after overfitting 16 examples: 0.0001

Rather write it yourself?

In the lesson you write this from a starter, in your browser with nothing to install, and checks tell you what's wrong as you go. A free account opens it and all 244 lessons.

Create a free account

Step 4: Get a baseline to beat

Fit a simple model on exactly the same data. Now the network has a number to beat.

baseline = LogisticRegression(C=0.1, max_iter=2000).fit(X_train.numpy(), y_train.numpy())
print(f"logistic regression validation AUC: {roc_auc_score(y_val.numpy(), baseline.predict_proba(X_val.numpy())[:, 1]):.3f}")
What it prints
logistic regression validation AUC: 0.853

Step 5: Train with early stopping

Dropout and weight decay fight overfitting. After every epoch, measure AUC on the validation set and keep a copy of the best weights. When 20 epochs pass with no improvement, stop and go back to that copy. Note model.eval() inside churn_prob: without it, dropout would add noise to your scores.

def churn_prob(model, X):
    model.eval()
    with torch.no_grad():
        return torch.softmax(model(X), dim=1)[:, 1].numpy()

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(X_train.shape[1], 32), nn.ReLU(), nn.Dropout(0.2),
                      nn.Linear(32, 32), nn.ReLU(), nn.Dropout(0.2), nn.Linear(32, 2))
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
best_auc, best_epoch, best_state = -1.0, 0, None
for epoch in range(150):
    model.train()
    order = torch.randperm(len(X_train))
    for i in range(0, len(X_train), 32):
        idx = order[i:i + 32]
        opt.zero_grad()
        nn.functional.cross_entropy(model(X_train[idx]), y_train[idx]).backward()
        opt.step()
    auc = roc_auc_score(y_val.numpy(), churn_prob(model, X_val))
    if auc > best_auc:
        best_auc, best_epoch = auc, epoch
        best_state = {k: v.clone() for k, v in model.state_dict().items()}
    if epoch - best_epoch >= 20:                    # no improvement for 20 epochs: stop
        break
model.load_state_dict(best_state)
print(f"network: best validation AUC {best_auc:.3f} at epoch {best_epoch}, stopped after epoch {epoch}")
What it prints
network: best validation AUC 0.864 at epoch 3, stopped after epoch 23

Step 6: Score each model once

The network looked better on validation data and was worse on the test set. With 600 rows of tabular data, that's common. The honest conclusion is to ship the simple model, which is cheaper to run and easier to explain.

baseline = LogisticRegression(C=0.1, max_iter=2000).fit(X_train.numpy(), y_train.numpy())
print("\ntest AUC (each model scored once):")
print(f"  logistic regression  {roc_auc_score(y_test.numpy(), baseline.predict_proba(X_test.numpy())[:, 1]):.3f}")
print(f"  neural network       {roc_auc_score(y_test.numpy(), churn_prob(model, X_test)):.3f}")
What it prints
test AUC (each model scored once):
  logistic regression  0.802
  neural network       0.789

A network earns its cost when there's a lot of data, or when the data has structure such as text or images. Here, the check you just ran saved a team from shipping something worse.

Step 7: Make it a reusable function

The same training loop, wrapped so you can call it again with different settings. This is what you write in the lesson's exercise, where checks confirm that it really stops early and returns the best weights.

def train_mlp(X_train, y_train, X_val, y_val, hidden=32, epochs=150, patience=20, lr=3e-3):
    model = nn.Sequential(nn.Linear(X_train.shape[1], hidden), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden, 2))
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    best_auc, best_epoch, best_state = -1.0, 0, None
    for epoch in range(epochs):
        model.train()
        order = torch.randperm(len(X_train))
        for i in range(0, len(X_train), 32):
            idx = order[i:i + 32]
            opt.zero_grad()
            nn.functional.cross_entropy(model(X_train[idx]), y_train[idx]).backward()
            opt.step()
        model.eval()
        with torch.no_grad():
            probs = torch.softmax(model(X_val), dim=1)[:, 1].numpy()
        auc = roc_auc_score(y_val.numpy(), probs)
        if auc > best_auc:
            best_auc, best_epoch = auc, epoch
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
        if epoch - best_epoch >= patience:
            break
    model.load_state_dict(best_state)
    model.eval()
    return model, float(best_auc), best_epoch

Where to take it next

  • Try more data: a network's advantage usually shows up as the dataset grows.
  • Search over hidden and lr on the validation set, and keep the test set for the end.
  • Add calibration if the scores will drive a decision directly.

Mistakes to watch for

  • Assuming a neural network beats simpler models on any dataset.
  • Fitting preprocessing on all the data.
  • Choosing the model by its test score after many tries.
  • Treating 'no improvement' as a failed project rather than a useful answer.
The complete program
import numpy as np
import pandas as pd
import requests
import torch
import torch.nn as nn
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

df = pd.DataFrame(requests.get("https://api.gradient.test/v1/customers").json()["data"])
numeric = ["tenure_months", "orders_last_year", "avg_order_value", "days_since_last_order",
           "discount_share", "email_open_rate", "support_tickets", "trail_club"]
categorical = ["region", "channel", "top_category"]
df["email_missing"] = df["email_open_rate"].isna().astype(float)
dev, test = train_test_split(df, test_size=0.2, random_state=7, stratify=df["churned"])
train, val = train_test_split(dev, test_size=0.25, random_state=7, stratify=dev["churned"])

def prepare(fit, *frames):
    med = fit[numeric].median()
    mu, sd = fit[numeric].fillna(med).mean(), fit[numeric].fillna(med).std()
    cats = pd.get_dummies(fit[categorical].replace("", "unknown")).columns
    out = []
    for f in frames:
        x = (f[numeric].fillna(med) - mu) / sd
        d = pd.get_dummies(f[categorical].replace("", "unknown")).reindex(columns=cats, fill_value=0)
        out.append(torch.tensor(np.c_[x.to_numpy(float), f[["email_missing"]].to_numpy(float), d.to_numpy(float)], dtype=torch.float32))
    return out

X_train, X_val, X_test = prepare(train, train, val, test)
y_train, y_val, y_test = (torch.tensor(f["churned"].to_numpy()) for f in (train, val, test))

def churn_prob(model, X):
    model.eval()
    with torch.no_grad():
        return torch.softmax(model(X), dim=1)[:, 1].numpy()

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(X_train.shape[1], 32), nn.ReLU(), nn.Dropout(0.2),
                      nn.Linear(32, 32), nn.ReLU(), nn.Dropout(0.2), nn.Linear(32, 2))
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
best_auc, best_epoch, best_state = -1.0, 0, None
for epoch in range(150):
    model.train()
    order = torch.randperm(len(X_train))
    for i in range(0, len(X_train), 32):
        idx = order[i:i + 32]
        opt.zero_grad()
        nn.functional.cross_entropy(model(X_train[idx]), y_train[idx]).backward()
        opt.step()
    auc = roc_auc_score(y_val.numpy(), churn_prob(model, X_val))
    if auc > best_auc:
        best_auc, best_epoch = auc, epoch
        best_state = {k: v.clone() for k, v in model.state_dict().items()}
    if epoch - best_epoch >= 20:                    # no improvement for 20 epochs: stop
        break
model.load_state_dict(best_state)
print(f"network: best validation AUC {best_auc:.3f} at epoch {best_epoch}, stopped after epoch {epoch}")

baseline = LogisticRegression(C=0.1, max_iter=2000).fit(X_train.numpy(), y_train.numpy())
print("\ntest AUC (each model scored once):")
print(f"  logistic regression  {roc_auc_score(y_test.numpy(), baseline.predict_proba(X_test.numpy())[:, 1]):.3f}")
print(f"  neural network       {roc_auc_score(y_test.numpy(), churn_prob(model, X_test)):.3f}")

Write it yourself

In the lesson you write this from a starter, in your browser with nothing to install, and checks tell you what's wrong as you go. A free account opens it and all 244 lessons.

Create a free account