5.3.3 PyTorch Pipeline for a Simple Machine Learning Problem

Why a full pipeline?

In A Simple MLP in PyTorch, the data came ready-made from scikit-learn, and we trained one network carefully. A contest task arrives differently: a statement that names the target, the metric and the submission format, a training file with labels, and a test file without them. The pipeline is the whole route from those files to a submission file.

Most lost points in contests come from that route, not from the model: a validation score that is too good to be true because information leaked from the validation rows, a test set scaled differently from the training set, or a submission with the wrong columns or rows in the wrong order. This page goes through the route once, step by step, with PyTorch, on a small medical dataset. The same steps work for any tabular task. Approaching a Typical Problem covers the same route with scikit-learn models, and the Titanic Walkthrough is a longer practice run of it.

What you need to know

By the end of this module, you should be comfortable with these ideas:

Learn it properly

This page is a worked example. The course below shows the same workflow on an even simpler problem, with more explanation of each step.

Pick one main course

Our pick: Learn PyTorch for Deep Learning, chapter 01

Read 01. PyTorch Workflow Fundamentals from Daniel Bourke's free course, if you have not already done it for PyTorch Basics. It goes through the same stages as this page, from preparing data to saving and loading a trained model, on a straight-line problem where every step is easy to see. Then come back here for the version with real files.

If you prefer something else

• The official tutorials: Optimizing Model Parameters and Save and Load the Model cover the training and saving stages.
• For the scaling step: scikit-learn's guide to standardization explains what StandardScaler does.

A quick tour

The rest of this page solves one task from start to finish. Put the files in the same folder as a new notebook, and run each cell as you read.

The task

Download the four files: train.csv, test.csv, sample_submission.csv and test_answers.csv. The last one holds the true test labels. In a real contest only the graders have it; here it lets you score yourself at the very end, so do not open it before then.

The data is the Breast Cancer Wisconsin (Diagnostic) dataset from the UCI Machine Learning Repository (CC BY 4.0), as shipped with scikit-learn. We split it into these files for this page.

Before any code, note the four facts a contestant should take from a statement:

• Target: malignant, either 0 or 1.
• Problem type: binary classification.
• Metric: accuracy, so the submission holds hard labels (0 or 1), not probabilities.
• Submission format: the columns id and malignant, one row per test row.

Here is the whole route at a glance, before we build it step by step:

1. Load and inspect the files

import pandas as pd

train = pd.read_csv("train.csv")
test = pd.read_csv("test.csv")
train.shape, test.shape

((455, 32), (114, 31))

The training file has one more column than the test file: the label. Let us count the two classes:

train["malignant"].value_counts()

malignant
0    285
1    170
Name: count, dtype: int64

285 benign and 170 malignant tumours, so about 63% of the rows are benign. A model that always answers 0 would already be right about 63% of the time: that is the number any real model has to beat.

Now look at the scale of a few features:

columns = ["mean_radius", "mean_area", "mean_smoothness"]
train[columns].describe().loc[["min", "mean", "max"]].round(3)

The features live on completely different scales: mean_area runs into the thousands, while mean_smoothness stays below 0.2. A network trained on the raw numbers would be dominated by the large columns, as Gradient Descent explains, so we will scale them in step 3.

2. Split off a validation set

The test file has no labels, so we cannot measure anything on it. We set aside 20% of the labelled rows as a validation set, keeping the two classes in the same proportions:

from sklearn.model_selection import train_test_split

train_part, val_part = train_test_split(
    train, test_size=0.2, stratify=train["malignant"],
    random_state=0)
len(train_part), len(val_part)

(364, 91)

364 rows to train on and 91 to check with. Every decision from now on, such as which epoch to keep, is made on the validation rows. The test rows are only used to make the final predictions.

3. Scale the features, without leakage

The feature columns are every column except id and the label. The id is just a row number, and a model must never see it: it carries no information about the tumour, and if the rows happen to be ordered, it can leak the answer.

feature_columns = train.columns.drop(["id", "malignant"])
len(feature_columns)

30

A StandardScaler shifts and stretches each column so that it has mean 0 and standard deviation 1. It learns each column's mean and standard deviation in fit, and applies them in transform:

from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()
scaler.fit(train_part[feature_columns])     # learn from training rows

X_train = scaler.transform(train_part[feature_columns])
X_val = scaler.transform(val_part[feature_columns])
X_test = scaler.transform(test[feature_columns])
y_train = train_part["malignant"].values
y_val = val_part["malignant"].values
X_train.shape, X_val.shape, X_test.shape

((364, 30), (91, 30), (114, 30))

The scaler is fitted on the training rows only. If it also saw the validation rows, their means and spreads would shape the scaling, and the validation score would no longer measure how the model does on truly new data. This is leakage. The same fitted scaler then transforms all three parts, so that a value of, say, 1.5 means the same thing in every one of them.

4. Tensors and DataLoaders

The rest is the routine from Dataset and DataLoader. For a yes/no label with BCEWithLogitsLoss, the labels are float32, not long:

import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

def make_dataset(features, labels):
    features = torch.tensor(features, dtype=torch.float32)
    labels = torch.tensor(labels, dtype=torch.float32)
    return TensorDataset(features, labels)

train_loader = DataLoader(make_dataset(X_train, y_train),
                          batch_size=32, shuffle=True)
val_loader = DataLoader(make_dataset(X_val, y_val), batch_size=256)

5. The model

One hidden layer of 64 neurons is plenty for 30 features and a few hundred rows. The last layer has a single output: one logit per row, where a positive logit means "malignant" and a negative one means "benign".

class TumourMLP(nn.Module):
    def __init__(self, n_features, n_hidden):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(n_features, n_hidden),
            nn.ReLU(),
            nn.Linear(n_hidden, 1),
        )

    def forward(self, x):
        return self.layers(x).squeeze(1)    # (batch, 1) -> (batch,)

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)
model = TumourMLP(n_features=30, n_hidden=64).to(device)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

The .squeeze(1) inside forward turns the output of shape (batch, 1) into (batch,), the same shape as the labels, which BCEWithLogitsLoss requires. Doing it once in the model means we cannot forget it anywhere else. Loss Functions explains the loss itself.

Check the starting loss before training, as in A Simple MLP in PyTorch. With two classes, an untrained model should give a probability of about 0.5 to every row, so the loss should be near $\ln 2 \approx 0.69$:

x_batch, y_batch = next(iter(train_loader))
logits = model(x_batch.to(device))
round(loss_fn(logits, y_batch.to(device)).item(), 3)

0.665

6. Train, and keep the best weights

The two functions are the ones from the previous module, with one change for a yes/no label: a row is predicted "malignant" when its logit is above 0, which is the same as a probability above 0.5.

def train_one_epoch(model, loader, loss_fn, optimizer):
    model.train()
    total_loss = 0.0
    for x_batch, y_batch in loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)
        loss = loss_fn(model(x_batch), y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(y_batch)
    return total_loss / len(loader.dataset)

def evaluate(model, loader, loss_fn):
    model.eval()
    total_loss = 0.0
    correct = 0
    with torch.no_grad():
        for x_batch, y_batch in loader:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)
            logits = model(x_batch)
            total_loss += loss_fn(logits, y_batch).item() * len(y_batch)
            predictions = (logits > 0).float()
            correct += (predictions == y_batch).sum().item()
    n_samples = len(loader.dataset)
    return total_loss / n_samples, correct / n_samples

import copy

best_val_loss = float("inf")
for epoch in range(1, 101):
    train_loss = train_one_epoch(model, train_loader, loss_fn,
                                 optimizer)
    val_loss, val_accuracy = evaluate(model, val_loader, loss_fn)
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_epoch = epoch
        best_accuracy = val_accuracy
        best_state = copy.deepcopy(model.state_dict())
    if epoch % 20 == 0:
        print(epoch, round(train_loss, 3), round(val_loss, 3),
              round(val_accuracy, 3))

model.load_state_dict(best_state)
torch.save(best_state, "tumour_mlp.pt")
print("best epoch:", best_epoch, "validation accuracy:",
      round(best_accuracy, 3))

20 0.07 0.066 0.978
40 0.042 0.054 0.967
60 0.028 0.057 0.967
80 0.02 0.056 0.978
100 0.014 0.059 0.978
best epoch: 39 validation accuracy: 0.967

The validation loss is lowest at epoch 39, where the model gets 88 of the 91 validation rows right. After that, the training loss keeps falling while the validation loss creeps up a little: the same overfitting pattern as in A Simple MLP in PyTorch. So we load the weights from epoch 39, and also save them to a file.

Notice that the accuracy at epochs 80 and 100 is higher, 0.978. That is exactly one more row out of 91. On a validation set this small, one row is noise, and the loss, which looks at every probability rather than just right or wrong, is the steadier guide for choosing an epoch.

7. Compare with a baseline

A score means little on its own. The simplest baseline predicts the more common class for every row; a stronger one is logistic regression on the same scaled features:

from sklearn.linear_model import LogisticRegression

majority = (y_val == 0).mean()
baseline = LogisticRegression(max_iter=1000)
baseline.fit(X_train, y_train)
print("majority:", round(majority, 3))
print("logistic regression:", round(baseline.score(X_val, y_val), 3))

majority: 0.626
logistic regression: 0.978

Both models are far above the majority guess. Logistic regression gets 89 of the 91 validation rows right and the network 88: one row apart, which at this size is noise. On a small table with well-scaled features, a linear model is a strong competitor, and that is worth knowing in a contest: it trains in seconds, so submit it first. We continue with the network, since this page is about PyTorch.

8. Predict the test rows

The fitted scaler has already transformed the test rows into X_test. The model, with the best weights loaded, turns them into logits, and a logit above 0 becomes the label 1:

model.eval()
with torch.no_grad():
    test_features = torch.tensor(X_test, dtype=torch.float32)
    test_logits = model(test_features.to(device))
probabilities = torch.sigmoid(test_logits).cpu()
predictions = (probabilities > 0.5).int()
print(probabilities[:5].round(decimals=3))
print(predictions[:5])

tensor([1.0000, 0.8580, 0.9590, 0.9410, 0.0010])
tensor([1, 1, 1, 1, 0], dtype=torch.int32)

torch.sigmoid turns each logit into a probability of "malignant". .cpu() brings the result back from the GPU, if there is one, so that pandas can use it. The task is scored by accuracy, so we submit the labels; a task scored on probabilities, such as by ROC AUC, would take probabilities instead (see Model Evaluation Metrics).

9. Write the submission, and check it

The rows of the submission must match the rows of test.csv, so we take the ids straight from it:

submission = pd.DataFrame({
    "id": test["id"],
    "malignant": predictions.numpy(),
})
submission.to_csv("submission.csv", index=False)

index=False stops pandas from writing its row numbers as an extra first column. Now read the file back, exactly as the graders will, and check it against the sample submission:

checked = pd.read_csv("submission.csv")
sample = pd.read_csv("sample_submission.csv")

assert list(checked.columns) == list(sample.columns)
assert len(checked) == len(test)
assert (checked["id"] == test["id"]).all()
assert checked["malignant"].isin([0, 1]).all()
checked.head()

No error, so all four checks passed.

Each assert stops the notebook with an error if its check fails, so a silent mistake cannot reach the graders. These four checks catch the most common submission errors: wrong columns, missing or extra rows, shuffled ids, and values of the wrong kind.

10. Score it yourself

In a contest, this is where you would submit. Here, we can play the grader:

from sklearn.metrics import accuracy_score

answers = pd.read_csv("test_answers.csv")
score = accuracy_score(answers["malignant"], checked["malignant"])
round(score, 3)

0.982

112 of the 114 test tumours are right, close to the validation accuracy, which is what an honest validation set promises. The logistic regression baseline scores 0.965 on the same test rows, 110 of 114: again within two rows of the network. On data like this, both are good answers.

Keeping the notebook contest-ready

• Run it from top to bottom. Before you submit, restart the kernel and run every cell in order. A notebook that only works because of cells run out of order will fail when it matters.
• Fix the seeds. random_state in the split and torch.manual_seed before the model make every run give the same numbers, so a change in score comes from your change, not from luck.
• Save the best weights to a file, as step 6 does, so a crash does not cost you the training time.
• Submit a baseline early. Logistic regression takes seconds. Submitting it first guarantees a score while you work on the network.

Common mistakes

How to practise

There is no PyTorch section on the practice sites we use (CodeChef, CS50 and HackerRank), so practise on real tasks and experiments:

• Remove the scaler (use the raw columns as the features) and train again. What happens to the curves and the validation accuracy?
• Build the same pipeline for the Titanic data from the Titanic Walkthrough. Its text columns need encoding first, as that page shows. Compare the network with the walkthrough's scikit-learn models.
• Change one thing at a time, such as the hidden size, the learning rate or the number of epochs, and keep the validation split fixed so the comparisons are fair.

Try XOR yourself

Here is a small test of whether you can adapt the pipeline. XOR has two inputs: the answer is 1 when exactly one input is 1, and 0 when the inputs match. A straight line cannot separate its two classes, so a network needs a hidden layer and an activation such as nn.ReLU() to learn it. We met this idea in Multi-Layer Perceptrons.

Start with these four rows. This tiny exercise runs fine on a CPU. Try building and training the model yourself before opening the hints. Your check is simple: can it predict all four answers?

import torch

X_xor = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y_xor = torch.tensor([[0.], [1.], [1.], [0.]])

Check yourself

What comes next

• Simple Problem Examples: the same pipeline for regression, several classes and several labels at once, and what changes each time.
• Model Evaluation Metrics: what to do when a task is scored by something other than accuracy.
• Titanic Walkthrough: a longer practice run of the whole route, with messier data.