Table of Contents
Problem Setup and the Classification Goal
A classification model predicts a discrete class label from input features. In the simplest case, binary classification, the label is either class 0 or class 1. Your dataset is a set of input vectors $x \in \mathbb{R}^d$ and labels $y \in \{0, 1\}$. The model produces a single real number called a logit, which is then turned into a probability for class 1.
In PyTorch, the most common beginner friendly setup for binary classification is to build a small neural network that outputs one logit per example and to train it with BCEWithLogitsLoss, which combines a numerically stable sigmoid plus binary cross entropy in one function.
A Minimal Synthetic Dataset for Classification
To focus on the end to end mechanics, it helps to start with a toy dataset where classes are separable enough to learn. The example below creates two clusters in 2D and labels them 0 and 1.
import torch
torch.manual_seed(0)
n_per_class = 500
x0 = torch.randn(n_per_class, 2) * 0.7 + torch.tensor([-2.0, 0.0])
x1 = torch.randn(n_per_class, 2) * 0.7 + torch.tensor([ 2.0, 0.0])
X = torch.cat([x0, x1], dim=0) # (1000, 2)
y = torch.cat([torch.zeros(n_per_class),
torch.ones(n_per_class)], dim=0) # (1000,)
BCEWithLogitsLoss expects floating targets for binary classification, so y should be float32. If you create labels as integers, convert them with y = y.float().
For binary classification with nn.BCEWithLogitsLoss, the model output must be raw logits of shape (N,) or (N, 1), and the target must be float with the same shape. Do not apply sigmoid in your model when using BCEWithLogitsLoss.
The Model: From Features to a Logit
A basic classification network is a small multilayer perceptron. It maps input features to a single logit.
import torch.nn as nn
model = nn.Sequential(
nn.Linear(2, 16),
nn.ReLU(),
nn.Linear(16, 1) # one logit
)
This model returns shape (N, 1) for a batch of size N. For convenience, you can remove the trailing dimension before computing loss and metrics.
Loss Function and Optimizer for Classification
For binary classification, use BCEWithLogitsLoss. It implements the binary cross entropy loss
$$
\mathcal{L} = -\frac{1}{N}\sum_{i=1}^N \left[y_i \log(\sigma(z_i)) + (1-y_i)\log(1-\sigma(z_i))\right]
$$
where $z_i$ is the logit and $\sigma$ is the sigmoid function.
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
If you use nn.CrossEntropyLoss for multi class classification, you must not use it for binary classification with a single output logit. For the single logit binary case, prefer BCEWithLogitsLoss.
Training: Forward, Loss, Backward, Step
This is a complete minimal training loop over the whole dataset as one batch. Later chapters will use DataLoaders and validation splits, but the basic mechanics are identical.
X_train = X
y_train = y.float()
for epoch in range(200):
model.train()
logits = model(X_train).squeeze(1) # (N,)
loss = loss_fn(logits, y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (epoch + 1) % 50 == 0:
print(f"epoch {epoch+1:3d} | loss {loss.item():.4f}")Turning Logits into Predictions and Accuracy
Logits are unbounded real numbers. For binary classification, convert logits to probabilities with sigmoid, then convert probabilities to class predictions with a threshold, typically 0.5.
$$
p = \sigma(z) = \frac{1}{1 + e^{-z}}, \quad \hat{y} = \mathbb{1}[p \ge 0.5]
$$
model.eval()
with torch.no_grad():
logits = model(X_train).squeeze(1)
probs = torch.sigmoid(logits)
preds = (probs >= 0.5).float()
acc = (preds == y_train).float().mean()
print("train accuracy:", acc.item())
Compute metrics from probabilities or predicted classes, but compute the loss from logits when using BCEWithLogitsLoss. Mixing these up is a common source of slow or unstable training.
A Quick Multi Class Variant (When You Have More Than Two Classes)
If you have $K$ classes, the common pattern is to output $K$ logits per example and train with nn.CrossEntropyLoss. Labels must be integer class indices, not one hot vectors.
K = 3
model_mc = nn.Sequential(
nn.Linear(10, 32),
nn.ReLU(),
nn.Linear(32, K) # K logits
)
loss_fn_mc = nn.CrossEntropyLoss()
X_mc = torch.randn(128, 10)
y_mc = torch.randint(0, K, (128,)) # int64 class indices
logits_mc = model_mc(X_mc) # (128, K)
loss_mc = loss_fn_mc(logits_mc, y_mc)
With nn.CrossEntropyLoss, pass raw logits of shape (N, K) and integer labels of shape (N,) with dtype torch.long. Do not apply softmax before the loss.
What “Good Output” Looks Like
In a working binary classifier, the average loss should decrease over epochs. Probabilities for class 1 should cluster near 0 for true class 0 examples and near 1 for true class 1 examples. Accuracy should rise above random chance, which is 50 percent for balanced binary classes. If accuracy stays near chance, the first things to check are label dtype and shape, whether you accidentally applied sigmoid twice, and whether your model output shape matches the target shape exactly.