Classification - Logits and Cross-Entropy

Context

A classifier is Lesson 4.3.1's network with its target swapped: a label, a 0 or a 1 per listing, in place of the price. The one output unit then returns a raw number, a function squashes that number into a probability, and a loss built for probabilities replaces nn.MSELoss, with the five steps of the loop untouched.

Every network that answers yes or no is built this way: whether a listing sells within 30 days, Lesson 3.4.1's question, whether a payment is fraud, whether a scan shows a tumor. Every network that picks one of many classes, Section 4.5's ten kinds of clothing, has one output unit per class and the loss of the last subtitle. Cross-entropy, the loss swapped in here, is the loss of every network classifier you will train from here on, the one object Lesson 4.2.2 said would be replaced.

The costliest mistakes in the swap do not stop the program. The loss takes probabilities where it wanted the raw numbers and returns a wrong number without a word, a threshold of 0.5 on probabilities that never reach 0.4 calls nothing sold, and the accuracy that follows, 0.798, reads as a model when it is the share of listings that did not sell. Lesson 4.3.1's network beat LinearRegression on the price by 7,140 of MAE. This lesson puts the same network on the label, shows what every piece returns, and reports what the runs show: on this label its best runs pass Lesson 3.4.1's logistic regression by about a hundredth of AUC.

Goal

• Build the label as a float32 (n, 1) tensor of zeros and ones beside Lesson 4.3.1's table with the standardized price as a 19th column, and state the majority baseline and LogisticRegression's AUC as the numbers to beat.
• Read a logit off a seeded network, turn it into a probability with torch.sigmoid and back with the log of the odds, and say where the 0.5 threshold sits on the logit scale.
• Compute nn.BCEWithLogitsLoss on logits and match it to the formula by hand, name the number it returns when handed probabilities, and diagnose its three refusals, an integer, a bool and a (n,) target.
• Train the classifier with the loss swapped into Lesson 4.2.3's loop, read its loss, AUC and confusion matrix at several thresholds, and report its AUC beside logistic regression's as the runs show it.
• Write the two-class form with a (n, 2) output, torch.softmax, argmax and nn.CrossEntropyLoss on int64 class indices, and show the number a softmax applied before that loss returns.

Prerequisites

None of this is required, the lesson stands on its own. Both are here if you want the loss drawn out before you compute it from a logit.

• YouTube - "Neural Networks Part 6: Cross Entropy" by StatQuest with Josh Starmer. Nine and a half minutes: cross-entropy defined on a softmax output, its general equation, the total over the rows of a table, and from 5:41 why a classifier is scored with it and not with the sum of squared residuals.
• Official docs - PyTorch 2.14, torch.nn.BCEWithLogitsLoss. The one class that combines the sigmoid and the binary cross-entropy, with the sentence on why that is more numerically stable than the two applied in turn, the formula per row, the shape rule that the target has the input's shape, and pos_weight, with its example of 100 positive rows against 300 negative ones.

Table of Contents

• A Label, Not a Price
• The Logit and the Sigmoid
• Binary Cross-Entropy, on Logits
• Training the Classifier, and the Threshold
• Many Classes: Softmax and CrossEntropyLoss

Theory

Every block in this lesson runs in ~/dev/bootcamp on the 2,875 listings with an area, split as Lesson 4.2.3 split them into 1,725 fit, 575 validation and 575 test rows, through the table Lesson 4.3.1's prep_net.py builds, with the standardized price appended as a 19th column. The loop is Lesson 4.2.3's, optim.SGD at a learning rate of 0.1 on the full batch, the five steps in Lesson 4.2.2's order with optimizer.zero_grad() last, and one object in it swapped, so it is written here and not imported from train_loop.py, whose run fits a price under nn.MSELoss. The first subtitle writes sold_prep.py on top of prep_net.py, imported from then on as 4.2.3's scripts imported prep.py. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

A Label, Not a Price

What. The target is sold_in_30_days, a bool column, True for a listing that sold within 30 days, built as a float32 tensor of zeros and ones in the output's shape, (n, 1). The input is Lesson 4.3.1's table with one column more, the standardized price appended last, 19 columns in all, the price an input as Section 3.4 made it. Why. A listing's price is a fair input to a model of whether it sells and never a feature of a model of the price, so it joins the table here and nowhere in Sections 4.2 and 4.4. The label is a float because the loss of the third subtitle multiplies with it, y * log(p), and the numbers to beat come first, because an accuracy on this label means nothing until the majority baseline sits beside it, Lesson 3.4.2's rule. When. Every time the target is a class and not a quantity, from here to the end of Module 6.

The table is Lesson 4.3.1's prep_net, imported, with the fit rows' standardized price appended as the 19th column, and the label is built beside it. Save it as sold_prep.py beside prep_net.py, because every later script of this lesson and Lesson 4.3.4's Lab 2 import it, and run it once, python3 sold_prep.py, for the lines under its main guard:

import torch
from prep_net import fit_rows, val_rows, test_rows, prep

# the fit rows' price statistics, applied to every set
price_mu, price_sd = fit_rows["price"].mean(), fit_rows["price"].std(ddof=0)

def features(frame):
    X, _ = prep(frame)
    z = (frame["price"] - price_mu) / price_sd
    price = torch.tensor(z.to_numpy(), dtype=torch.float32).unsqueeze(1)
    return torch.cat([X, price], dim=1)

def label(frame):
    sold = torch.tensor(frame["sold_in_30_days"].to_numpy(), dtype=torch.float32)
    return sold.unsqueeze(1)

X_fit, y_fit = features(fit_rows), label(fit_rows)
X_val, y_val = features(val_rows), label(val_rows)
X_test, y_test = features(test_rows), label(test_rows)

if __name__ == "__main__":
    print(X_fit.shape, X_val.shape, X_test.shape)
    print(X_fit[0])
    print(fit_rows["sold_in_30_days"].head(3).tolist(), fit_rows["sold_in_30_days"].dtype)
    print(y_fit[:3].squeeze(1), y_fit.dtype, y_fit.shape)
    print(round(y_fit.mean().item(), 3), round(y_val.mean().item(), 3), int(y_val.sum()))

torch.Size([1725, 19]) torch.Size([575, 19]) torch.Size([575, 19])
tensor([-0.8938, -0.3373, -0.8639,  1.0243,  0.0000,  0.0000,  0.0000,  0.0000,
         0.0000,  0.0000,  1.0000,  1.0000,  0.0000,  0.0000,  1.0000,  0.0000,
         0.0000,  0.0000, -0.5459])
[True, True, False] bool
tensor([1., 1., 0.]) torch.float32 torch.Size([1725, 1])
0.209 0.202 116

The import of prep_net runs everything above its main guard, Lesson 4.2.3's pattern, so the frames, the statistics, the encoder and prep are Lesson 4.3.1's, and the fit rows are the 1,725 whose mean area is 99.44. features keeps the 18 columns prep returns, drops the price target beside them into _, Lesson 1.2.3's name for an ignored value, standardizes the price with the fit rows' mean and std(ddof=0) as prep standardizes the area, and torch.cat, Lesson 4.1.2's, joins it along dim=1 as the 19th and last column: -0.5459 for the first fit row, 272,000 against a mean of 384,795 and a standard deviation of 206,615, after the four numeric columns and the encoder's 14 in Lesson 4.3.1's order. The label arrives as True, True, False and leaves as 1., 1., 0., float32 and (1725, 1), the dtype and shape of the network's output. The fit rows' share sold is 0.209 and the validation rows' 0.202, 116 of 575, so a 0 for every validation row scores 0.798, the majority baseline Lesson 3.4.2 set beside every accuracy. Logistic regression on the same 19 columns gives the second number:

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, roc_auc_score
from sold_prep import X_fit, y_fit, X_val, y_val

# the labels as one-dimensional NumPy arrays, the form the scorers take
sold_fit = y_fit.squeeze(1).numpy()
sold_val = y_val.squeeze(1).numpy()
never = np.zeros(len(sold_val))
# 0.798, the majority baseline: no listing called sold
print(round(accuracy_score(sold_val, never), 3))
logreg = LogisticRegression(max_iter=2000).fit(X_fit.numpy(), sold_fit)
called = logreg.predict(X_val.numpy())
proba = logreg.predict_proba(X_val.numpy())[:, 1]
# 1 0.797, one listing called sold, and it did not sell
print(int(called.sum()), round(accuracy_score(sold_val, called), 3))
# 0.665, the number to beat
print(round(roc_auc_score(sold_val, proba), 3))

Logistic regression's one call puts its accuracy one listing under the baseline, Lesson 3.4.3's finding on a model whose probabilities seldom cross 0.5, and its AUC is Lesson 3.4.3's area under the ROC curve.

DON'T:

• DON'T hand the label to a loss as the bool it arrives as. torch.tensor(frame["sold_in_30_days"].to_numpy()) is a torch.bool tensor, and the third subtitle shows the refusal it meets.
• DON'T read 0.797 as a model that nearly works. 0.798 is what predicting that nothing sells scores, and every accuracy on this label is read beside it.
• DON'T put the price into a table whose target is the price. It is a fair input here because the target is the sale, and Lesson 4.3.1's table left it out for the same reason.

The Logit and the Sigmoid

What. A logit (the raw output of a classifier's last layer, any real number, positive where the model leans toward calling the listing sold and negative where it leans the other way) is what the one output unit returns, and the sigmoid, Lesson 3.4.1's logistic function 1 / (1 + exp(-z)) under PyTorch's name, turns it into a probability, torch.sigmoid. The two are a pair: log(p / (1 - p)), 3.4.1's log-odds, turns the probability back into the logit. Why. A linear layer returns any real number and a probability lives between 0 and 1, so one has to be mapped onto the other. Lesson 3.4.1's predict called a listing sold where a score was above 0 and said that is where its probability is above one half: that score is the logit, and 0 on the logit scale is 0.5 on the probability scale. The network returns the logit and not the probability because the loss of the next subtitle wants the logit. When. Every binary classifier built on a sigmoid, in every library: a score before the sigmoid is a logit, and a number between 0 and 1 after it is a probability.

import torch
from torch import nn
from sold_prep import X_fit, y_fit, X_val, y_val

torch.manual_seed(42)
model = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 1))
# 673, Lesson 4.3.1's 641 and one more weight on each of the 32 units
print(sum(p.numel() for p in model.parameters()))
with torch.no_grad():
    logits = model(X_val)
print(logits.shape, logits.dtype)  # torch.Size([575, 1]) torch.float32
z = logits[:4].squeeze(1)
print(z)                           # tensor([0.2124, 0.2181, 0.1610, 0.2051])
print(torch.sigmoid(z))            # tensor([0.5529, 0.5543, 0.5402, 0.5511])
print(1 / (1 + torch.exp(-z)))     # tensor([0.5529, 0.5543, 0.5402, 0.5511])

Before any step the first four logits are small and positive, and their probabilities sit just above one half.

print(torch.sigmoid(torch.tensor([-5.0, -1.0, 0.0, 1.0, 5.0])))
print(torch.log(torch.sigmoid(z) / (1 - torch.sigmoid(z))))
print(round(logits.min().item(), 4), round(logits.max().item(), 4))
print(torch.equal(torch.sigmoid(logits) >= 0.5, logits >= 0), int((logits >= 0).sum()))

tensor([0.0067, 0.2689, 0.5000, 0.7311, 0.9933])
tensor([0.2124, 0.2181, 0.1610, 0.2051])
-0.0095 0.4015
True 574

The sigmoid of 0 is exactly 0.5, of -5 it is 0.0067 and of 5 it is 0.9933, so a logit far from 0 is a probability near 0 or 1 and a logit near 0 is a coin. The log of the odds returns the four logits to the digit. The seeded network's logits run from -0.0095 to 0.4015, all but one above 0, so logits >= 0 and sigmoid(logits) >= 0.5 are the same 575 answers, 574 of them sold: the untrained network calls every validation listing but one sold, and a threshold of 0.5 on a probability is a threshold of 0 on a logit.

DON'T:

• DON'T read a logit as a probability. 0.2124 is not a 21 percent chance of a sale, it is a logit whose probability is 0.5529.
• DON'T put a sigmoid on the end of the network. The network returns logits so the loss can take them, and the probability is torch.sigmoid(model(X)) where a probability is read, in a validation pass or a report, and nowhere in the loss.

Binary Cross-Entropy, on Logits

What. nn.BCEWithLogitsLoss() is binary cross-entropy (Lesson 3.4.1's log loss under PyTorch's name) computed from the logits: it applies the sigmoid itself, inside, in one step. It takes (n, 1) logits and a float32 (n, 1) target of zeros and ones, and it refuses an integer target, a bool target and a (n,) target. Why. A squared error on a probability charges a confident wrong answer, 0.01 for a listing that sold, 0.9801 and never more than 1, and this loss charges it -log(0.01), 4.605, so a model that is sure and wrong pays most. Computing from the logit keeps the logarithm finite, shown below on a logit of -100, and the shape rule nn.MSELoss only warned about is enforced here with a ValueError. When. Every binary classifier's loss, with the logits going in and never the probabilities.

loss_fn = nn.BCEWithLogitsLoss()
start = loss_fn(logits, y_val)
print(f"{start.item():.4f}")  # 0.7530
p = torch.sigmoid(logits)
charge = -(y_val * torch.log(p) + (1 - y_val) * torch.log(1 - p))
print(y_val[:4].squeeze(1))   # tensor([0., 1., 0., 0.])
print(charge[:4].squeeze(1))  # tensor([0.8050, 0.5900, 0.7769, 0.8009])
# 0.7530 True
print(f"{charge.mean().item():.4f}", torch.allclose(charge.mean(), start))

The four charges show the rule row by row: the first listing did not sell and was given 0.5529, so it is charged -log(1 - 0.5529), 0.8050, the most of the four because it sits furthest from its label, and the second sold at 0.5543 and is charged -log(0.5543), 0.5900, the least. A probability near a row's label costs almost nothing, and one far from it costs without bound. The same loss handed the probabilities:

wrong = loss_fn(p, y_val)
print(f"{wrong.item():.4f}")            # 0.8928
print(torch.sigmoid(p[:4]).squeeze(1))  # tensor([0.6348, 0.6351, 0.6319, 0.6344])

0.8928, no warning, no error. The loss applied its sigmoid to numbers that had already been through one, so it scored 0.6348 for the first row where the model said 0.5529, and every row above 0.6 whatever the label, which is a different number minimized at a different place. Two baselines put the 0.7530 in place:

base_rate = y_fit.mean()
base_logit = torch.log(base_rate / (1 - base_rate))
# 0.209 -1.333
print(round(base_rate.item(), 3), round(base_logit.item(), 3))
flat = loss_fn(torch.full_like(y_val, base_logit.item()), y_val)
print(f"{flat.item():.4f}")  # 0.5030
coin = loss_fn(torch.zeros_like(y_val), y_val)
print(f"{coin.item():.4f}")  # 0.6931

The base rate's logit for every validation row is Lesson 3.4.1's flat quote, and its loss is the line a trained classifier has to cross, as Lesson 4.2.3's mean price was for the regression. A logit of 0 for every row, a coin, scores the natural logarithm of 2, and the seeded network's 0.7530 sits above both: all but one of its probabilities are above 0.5, so 458 of the 459 unsold rows are charged more than the coin's 0.6931. Why the loss takes the logit and does the sigmoid itself:

far, one = torch.tensor([[-100.0]]), torch.tensor([[1.0]])
print(torch.log(torch.sigmoid(far)))  # tensor([[-inf]])
print(loss_fn(far, one))              # tensor(100.)

A logit of -100 is a float32 probability of exactly 0, whose logarithm is -inf, the number the hand formula would average. nn.BCEWithLogitsLoss computes the charge from the logit and returns 100, the finite number the formula means, which is why the network ends on a logit and the loss carries the sigmoid. The three refusals:

try:
    loss_fn(logits, y_val.long())
except RuntimeError as err:
    print(err)
try:
    loss_fn(logits, y_val.bool())
except RuntimeError as err:
    print(err)
try:
    loss_fn(logits, y_val.squeeze(1))
except ValueError as err:
    print(err)

result type Float can't be cast to the desired output type Long
Subtraction, the `-` operator, with two bool tensors is not supported. Use the `^` or `logical_xor()` operator instead.
Target size (torch.Size([575])) must be the same as input size (torch.Size([575, 1]))

An int64 target is refused with a Float result that cannot be written into a Long, a bool target by the first subtraction it meets, with a message that names an operator and never the target, and a (n,) target against a (n, 1) output outright, a ValueError naming both sizes where nn.MSELoss warned and broadcast. The target this loss takes is the one label builds: float32, zeros and ones, (n, 1).

DON'T:

• DON'T apply torch.sigmoid before the loss. loss_fn(torch.sigmoid(logits), y) prints 0.8928 for the 0.7530 the logits score, raises nothing, and a loop on it minimizes a loss of probabilities put through the sigmoid twice.
• DON'T hand the loss a (n,) label for the (n, 1) output. The two shapes must match, and the refusal names both sizes when they differ.
• DON'T build the label as int64. It is refused as the bool is, and the dtype is set where the label is built, dtype=torch.float32.

Training the Classifier, and the Threshold

What. The loop is Lesson 4.2.3's with nn.MSELoss replaced by nn.BCEWithLogitsLoss and the price by the label, the same five steps and the same validation pass under both switches, with the probabilities read out of the pass through torch.sigmoid so that the AUC and the confusion matrix stand beside the loss. The threshold, Lesson 3.4.3's, is then a choice made on the probabilities the network returns, and 0.5 is one choice among them. Why. A classifier's loss falls and says nothing about what the model calls sold, and after 300 epochs at the section's rate the probabilities stay under 0.4, so 0.5 calls nothing and the accuracy is the majority baseline to the digit: the loss against its baseline, the AUC and a lower threshold are what show it learned anything. The comparison with Lesson 3.4.1's model is what the training is for. When. Every classifier, every run: the loss, the AUC and the confusion matrix at the threshold the question needs, never the accuracy at 0.5 alone.

import torch
from torch import nn, optim
from sklearn.metrics import accuracy_score, confusion_matrix, roc_auc_score
from sold_prep import X_fit, y_fit, X_val, y_val

loss_fn = nn.BCEWithLogitsLoss()
sold_val = y_val.squeeze(1).numpy()

def validate(model, X, y):
    model.eval()
    with torch.no_grad():
        out = model(X)
        loss = loss_fn(out, y)
    model.train()
    return loss.item(), torch.sigmoid(out).squeeze(1).numpy()

def train(lr, epochs, report=()):
    torch.manual_seed(42)
    model = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 1))
    optimizer = optim.SGD(model.parameters(), lr=lr)
    for epoch in range(1, epochs + 1):
        loss = loss_fn(model(X_fit), y_fit)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        val_loss, proba = validate(model, X_val, y_val)
        if epoch in report:
            auc = roc_auc_score(sold_val, proba)
            print(epoch, f"{loss.item():.4f}", f"{val_loss:.4f}", round(auc, 3))
    return proba

proba = train(0.1, 300, report=(1, 100, 300))

1 0.7526 0.7280 0.514
100 0.5029 0.4957 0.595
300 0.4955 0.4921 0.607

validate is Lesson 4.2.3's with one more return, the probabilities of the pass as a (575,) NumPy array, and train wraps the five steps so the loop can run four times from the same seed. Each reported line is the epoch, the training loss the step was taken from, the validation loss after it and the AUC of the validation probabilities, roc_auc_score as Lesson 3.4.3 called it. The AUC after the first step is 0.514, a ranking a hair better than a coin's, by epoch 100 the validation loss is 0.4957, under the base rate's 0.5030, and by 300 it is 0.4921 with an AUC of 0.607. What it calls sold is another matter:

print(round(proba.min(), 3), round(proba.max(), 3), round(proba.mean(), 3))
for threshold in (0.5, 0.3, 0.25, 0.2):
    called = proba >= threshold
    found = int((called & (sold_val == 1)).sum())
    accuracy = round(accuracy_score(sold_val, called), 3)
    print(threshold, int(called.sum()), found, accuracy, confusion_matrix(sold_val, called).ravel())

0.1 0.343 0.214
0.5 0 0 0.798 [459   0 116   0]
0.3 48 10 0.75 [421  38 106  10]
0.25 162 46 0.677 [343 116  70  46]
0.2 304 73 0.523 [228 231  43  73]

The probabilities run from 0.1 to 0.343 with a mean of 0.214, about the share that sold, and none reaches 0.5, so that threshold calls nothing sold: an accuracy of 0.798, the majority baseline to the digit, and a confusion matrix of 459 true negatives and 116 false negatives in ravel's order from Lesson 3.4.3. Each line is one threshold, the listings it calls, the sales among them, the accuracy and the four counts, and the accuracy falls at every step down, 0.750, 0.677, 0.523, because each lower threshold adds more unsold listings than sales, Lesson 3.4.3's trade on a network's probabilities. Which threshold is kept is chosen on validation rows by a rule written first, as 3.4.3 chose it, and nothing here chooses one. Whether more training changes the ranking, three more runs from the same seed:

for lr, epochs in ((0.1, 1000), (0.5, 300), (1.0, 300)):
    proba = train(lr, epochs, report=(epochs,))
    print(int((proba >= 0.5).sum()))

1000 0.4777 0.4816 0.646
0
300 0.4670 0.4763 0.659
0
300 0.4529 0.4782 0.659
1

A thousand epochs at 0.1 reach an AUC of 0.646, 300 at a rate of 0.5 reach 0.659 with a validation loss of 0.4763, and 300 at 1.0 reach 0.659 and call one listing sold at 0.5. Logistic regression on the same columns scored 0.665. The network's ranking rises toward that number with longer or faster training and, run longer still, passes it by about a hundredth: 2,000 epochs at 0.1 reach 0.672 and 500 at 0.5 reach 0.673. On this label, with these 19 columns, the hidden layer buys about a hundredth of AUC over Lesson 3.4.1's model, where on the price it bought 7,140 of MAE. That is a result, and it is reported as one.

DON'T:

• DON'T report the accuracy at 0.5. It printed 0.798 for a model that calls nothing sold and 0.798 for predicting that nothing sells, and the two are the same number for different reasons.
• DON'T take 0.5 as the threshold because predict used it in Module 3. The trained network's probabilities top out at 0.343, and every threshold that calls anything sold sits below that.
• DON'T expect a network to beat a linear model because it is a network. On the price it did, by 7,140 of MAE, and on this label its longer and faster runs reached 0.646 to 0.673 against logistic regression's 0.665.

Many Classes: Softmax and CrossEntropyLoss

What. A last layer of C units returns C logits per row, one per class, and the softmax (the function that turns a row of logits into probabilities that sum to 1, each logit's exponential divided by the row's sum of exponentials) makes the row's probabilities, torch.softmax(out, dim=1). The predicted class index (the integer naming a class, 0 or 1 here and 0 to 9 for the ten clothing classes of Section 4.5) is argmax(dim=1). nn.CrossEntropyLoss() is cross-entropy (the same log loss over C classes), computed from the (n, C) logits and a one-dimensional int64 target of class indices, with the softmax inside it, as the sigmoid is inside the binary loss. Why. The binary case is this with C of 2 and one logit kept, because the softmax of two logits [0, z] is the sigmoid of z, so the two forms are one loss. Section 4.5 has ten classes and no single logit can hold them, so its classifier takes this form, with its own two refusals and its own version of the activation applied twice. When. Whenever the target is one of several classes, and whenever you read a classifier someone else wrote: a (n, C) output with CrossEntropyLoss is this subtitle.

import torch
from torch import nn, optim
from sklearn.metrics import confusion_matrix, roc_auc_score
from sold_prep import X_fit, y_fit, X_val, y_val

torch.manual_seed(42)
# one more output unit than the classifier, its 32 weights and its bias
two = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 2))
print(sum(p.numel() for p in two.parameters()))
with torch.no_grad():
    out = two(X_val)
print(out.shape)
print(out[:2])
probs = torch.softmax(out, dim=1)
print(probs[:2])
print(probs[:2].sum(dim=1))
by_hand = torch.exp(out) / torch.exp(out).sum(dim=1, keepdim=True)
print(torch.allclose(by_hand, probs))
print(out[:2].argmax(dim=1))
print(torch.equal(out.argmax(dim=1), (probs[:, 1] >= 0.5).long()))

706
torch.Size([575, 2])
tensor([[0.0807, 0.0443],
        [0.0865, 0.1293]])
tensor([[0.5091, 0.4909],
        [0.4893, 0.5107]])
tensor([1., 1.])
True
tensor([0, 1])
True

The output is (575, 2), two logits per listing, class 0's, not sold, and class 1's, sold, in the order the class indices name them. torch.softmax along dim=1 works across the row and agrees with the formula by hand within allclose, so every row sums to 1. argmax(dim=1) returns the index of the larger, 0 for the first row and 1 for the second, and for two classes that is exactly probs[:, 1] >= 0.5 on all 575, so the binary form's 0.5 threshold is this form's argmax. The loss and its target:

classes = y_val.squeeze(1).long()
# torch.int64 torch.Size([575]) tensor([0, 1, 0])
print(classes.dtype, classes.shape, classes[:3])
ce = nn.CrossEntropyLoss()
right = ce(out, classes)
print(f"{right.item():.4f}")               # 0.7167
picked = probs[torch.arange(len(classes)), classes]
# 0.7167, minus the log of each row's own class's probability, averaged
print(f"{(-torch.log(picked)).mean().item():.4f}")
print(f"{ce(probs, classes).item():.4f}")  # 0.7046

The target is the label with its dimension removed and cast to int64, (575,), the class index of each row. probs[torch.arange(575), classes] indexes with two integer arrays, one per axis, and picks each row's probability for its own class, the probability the formula takes the log of. Handed the softmax's probabilities, the loss returns 0.7046 with no word said, the binary case's 0.8928 again: it applied its own softmax to rows that already summed to 1. The refusals:

try:
    ce(out, y_val.squeeze(1))
except RuntimeError as err:
    print(err)
try:
    ce(out, classes.unsqueeze(1))
except RuntimeError as err:
    print(err)

expected target dtype to be Long or Byte, but got Float
0D or 1D target tensor expected, multi-target not supported

The dtype rule is the binary loss's turned around: nn.BCEWithLogitsLoss refused the integer and took the float, nn.CrossEntropyLoss refuses the float label and takes the int64 one. The shape rule is its own: the target is one-dimensional, one index per row, and beside two logits the (n, 1) shape every regression target and the binary label took is refused as a multi-target. Why one logit was enough for two classes:

# tensor([[0.4471, 0.5529]])
print(torch.softmax(torch.tensor([[0.0, 0.2124]]), dim=1))
print(torch.sigmoid(torch.tensor(0.2124)))  # tensor(0.5529)

The one-unit form is the two-unit form with class 0's logit fixed at 0. Both forms train on the same loop with the loss and the target swapped:

sold_val = y_val.squeeze(1).numpy()
classes_fit = y_fit.squeeze(1).long()
torch.manual_seed(42)
two = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = optim.SGD(two.parameters(), lr=0.1)
for epoch in range(1, 301):
    loss = ce(two(X_fit), classes_fit)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    two.eval()
    with torch.no_grad():
        out = two(X_val)
        val_loss = ce(out, classes).item()
        proba = torch.softmax(out, dim=1)[:, 1].numpy()
        pred = out.argmax(dim=1).numpy()
    two.train()
    if epoch in (1, 100, 300):
        auc = roc_auc_score(sold_val, proba)
        print(epoch, f"{loss.item():.4f}", f"{val_loss:.4f}", round(auc, 3), int(pred.sum()))
print(confusion_matrix(sold_val, pred).ravel())

1 0.7178 0.6745 0.476 87
100 0.5016 0.4964 0.586 0
300 0.4923 0.4919 0.605 0
[459   0 116   0]

The probability of a sale is the softmax's second column and the call is argmax, which calls 87 listings sold after the first step and none at epochs 100 and 300. At 300 epochs the two-unit form reaches a validation loss of 0.4919 and an AUC of 0.605 beside the one-unit form's 0.4921 and 0.607, close and not equal, and its confusion matrix is the one-unit form's at 0.5. The two start from different draws and step differently: one step moves the gap between two logits twice as far as it moves one logit, and started from the one-unit form's own function with its last layer at half the rate, the two-unit form lands on 0.4921 and 0.607. The classifier Section 4.5 trains on ten classes ends in ten output units and takes this loss with a target of indices from 0 to 9.

DON'T:

• DON'T apply torch.softmax before nn.CrossEntropyLoss. It returned 0.7046 for the 0.7167 the logits score, said nothing, and the activation applied twice is the binary case's 0.8928 again.
• DON'T give nn.CrossEntropyLoss the float label or the (n, 1) one. Beside two logits both are refused, beside the one-unit form's single logit the float (n, 1) label is taken without a word and scores -0.0000 whatever the logit, and beside two logits y.squeeze(1).long() is the target it takes.
• DON'T read argmax as a choice of threshold. For two classes it is the 0.5 threshold on the second column, the 575 answers were equal, and a threshold the question needs is set on probs[:, 1].

Lab Exercises

Lab Exercise 1: One Neuron Is Logistic Regression

In neuron_start.py, import the table and the labels from sold_prep, seed with 42 and build nn.Linear(19, 1), one neuron on the 19 columns and nothing after it. Print the shapes of its weight and bias, then its first three weights and its bias to four decimals. Compute its logits for the validation rows under no_grad, print the first four and their probabilities through torch.sigmoid, and print whether torch.sigmoid and 1 / (1 + exp(-z)) agree on them under torch.allclose. Print the smallest and largest probability to three decimals beside how many listings the untrained neuron calls sold, counted on the logits.

Then print the start loss under nn.BCEWithLogitsLoss beside the formula written out by hand, to four decimals, the number the loss returns when handed the probabilities, and the AUC of the untrained neuron's probabilities to three decimals. Catch and print the refusal of an int64 label and of a (n,) label. Say in a closing comment why this model is Lesson 3.4.1's logistic regression and what the loss is in that lesson's words.

Lab Exercise 2: The Neuron Trained, in Both Forms

In neuron_trained.py, train the seeded nn.Linear(19, 1) with nn.BCEWithLogitsLoss at a learning rate of 0.1 for 300 epochs, with the validation pass under both switches, printing the training loss, the validation loss and the AUC at epochs 1, 100 and 300. Print the probabilities' smallest, largest and mean to three decimals, then for the thresholds 0.5, 0.3 and 0.2 the listings called sold, the sales among them, the accuracy and the confusion matrix's four counts.

Then write the same neuron as two logits, nn.Linear(19, 2) under the same seed, with nn.CrossEntropyLoss on int64 class indices, trained the same way, printing at the same epochs the two losses, the AUC of the softmax's second column and how many listings argmax calls sold. Print whether argmax and the second column at 0.5 agree on every validation row, and the confusion matrix of the argmax calls. Finish by training the one-logit neuron again for 3,000 epochs, printing its line at epochs 1,000 and 3,000, and set its last AUC beside LogisticRegression's 0.665 in a comment.

import torch
from torch import nn
from sklearn.metrics import roc_auc_score
from sold_prep import X_val, y_val

torch.manual_seed(42)
neuron = nn.Linear(19, 1)
print(neuron.weight.shape, neuron.bias.shape)
w = neuron.weight[0]
print(f"{w[0].item():.4f} {w[1].item():.4f} {w[2].item():.4f} {neuron.bias.item():.4f}")
with torch.no_grad():
    logits = neuron(X_val)
z = logits[:4].squeeze(1)
print(z)
print(torch.sigmoid(z))
print(torch.allclose(torch.sigmoid(z), 1 / (1 + torch.exp(-z))))
p = torch.sigmoid(logits)
print(round(p.min().item(), 3), round(p.max().item(), 3), int((logits >= 0).sum()))

loss_fn = nn.BCEWithLogitsLoss()
start = loss_fn(logits, y_val)
hand = -(y_val * torch.log(p) + (1 - y_val) * torch.log(1 - p)).mean()
print(f"{start.item():.4f} {hand.item():.4f}")
print(f"{loss_fn(p, y_val).item():.4f}")
print(round(roc_auc_score(y_val.squeeze(1).numpy(), p.squeeze(1).numpy()), 3))
try:
    loss_fn(logits, y_val.long())
except RuntimeError as err:
    print(err)
try:
    loss_fn(logits, y_val.squeeze(1))
except ValueError as err:
    print(err)
# One weight per column, one bias and a sigmoid on the score is Lesson 3.4.1's
# logistic regression, and this loss is its log loss: the fit is what differs.

torch.Size([1, 19]) torch.Size([1])
0.1754 0.1904 -0.0537 0.0585
tensor([0.8742, 0.4401, 0.2820, 1.4635])
tensor([0.7056, 0.6083, 0.5700, 0.8121])
True
0.338 0.848 333
0.8140 0.8140
0.9002
0.424
result type Float can't be cast to the desired output type Long
Target size (torch.Size([575])) must be the same as input size (torch.Size([575, 1]))

import numpy as np
import torch
from torch import nn, optim
from sklearn.metrics import accuracy_score, confusion_matrix, roc_auc_score
from sold_prep import X_fit, y_fit, X_val, y_val

sold_val = y_val.squeeze(1).numpy()
bce = nn.BCEWithLogitsLoss()

def train_one(epochs, report):
    torch.manual_seed(42)
    neuron = nn.Linear(19, 1)
    optimizer = optim.SGD(neuron.parameters(), lr=0.1)
    for epoch in range(1, epochs + 1):
        loss = bce(neuron(X_fit), y_fit)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        neuron.eval()
        with torch.no_grad():
            out = neuron(X_val)
            val_loss = bce(out, y_val).item()
            proba = torch.sigmoid(out).squeeze(1).numpy()
        neuron.train()
        if epoch in report:
            auc = roc_auc_score(sold_val, proba)
            print(epoch, f"{loss.item():.4f}", f"{val_loss:.4f}", round(auc, 3))
    return proba

proba = train_one(300, (1, 100, 300))
print(round(proba.min(), 3), round(proba.max(), 3), round(proba.mean(), 3))
for threshold in (0.5, 0.3, 0.2):
    called = proba >= threshold
    found = int((called & (sold_val == 1)).sum())
    accuracy = round(accuracy_score(sold_val, called), 3)
    print(threshold, int(called.sum()), found, accuracy, confusion_matrix(sold_val, called).ravel())

1 0.8172 0.7880 0.426
100 0.4997 0.4943 0.602
300 0.4933 0.4888 0.622
0.082 0.393 0.216
0.5 0 0 0.798 [459   0 116   0]
0.3 64 19 0.753 [414  45  97  19]
0.2 311 77 0.525 [225 234  39  77]

ce = nn.CrossEntropyLoss()
classes_fit, classes_val = y_fit.squeeze(1).long(), y_val.squeeze(1).long()
torch.manual_seed(42)
pair = nn.Linear(19, 2)
optimizer = optim.SGD(pair.parameters(), lr=0.1)
for epoch in range(1, 301):
    loss = ce(pair(X_fit), classes_fit)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    pair.eval()
    with torch.no_grad():
        out = pair(X_val)
        val_loss = ce(out, classes_val).item()
        proba_two = torch.softmax(out, dim=1)[:, 1].numpy()
        pred = out.argmax(dim=1).numpy()
    pair.train()
    if epoch in (1, 100, 300):
        auc = roc_auc_score(sold_val, proba_two)
        print(epoch, f"{loss.item():.4f}", f"{val_loss:.4f}", round(auc, 3), int(pred.sum()))
print(np.array_equal(pred, proba_two >= 0.5))
print(confusion_matrix(sold_val, pred).ravel())

long_run = train_one(3000, (1000, 3000))
# 0.662 at 3,000 epochs and still rising, against LogisticRegression's 0.665

1 0.6842 0.6652 0.559 251
100 0.5006 0.4969 0.583 0
300 0.4912 0.4893 0.618 0
True
[459   0 116   0]
1000 0.4848 0.4813 0.646
3000 0.4765 0.4746 0.662

Theory Assessment

Q1: What the Output Unit Returns

The seeded network returned 0.2124 for the first validation listing and torch.sigmoid turned it into 0.5529. Say what each number is, why the second sits just above one half, what log(p / (1 - p)) does to it, and where the threshold of 0.5 sits on the scale of the first number. Then say why the network returns the first number and not the second.

Hint: One of the two can be any real number and the other cannot. Ask what the loss of the third subtitle takes and what it does with it.

Q2: Three Refusals and One Silence

A teammate's loop passes torch.sigmoid(model(X_fit)) to nn.BCEWithLogitsLoss and prints a loss that falls. Say what the loss computes on those inputs and why nothing was raised, what it printed at the seeded start against the right number, and what the same loss does with an int64 label, a bool label and a (n,) label. Say which of the four cases is the one that runs to the end.

Hint: Count the sigmoids each row goes through. For the refusals, ask what the loss multiplies, what it subtracts, and what it compares before it computes anything.

Q3: The Model That Called Nothing Sold

After 300 epochs the network's accuracy at 0.5 was 0.798, the same as predicting that nothing sells, and a teammate concludes it learned nothing. Set the validation loss 0.4921 against 0.5030, the AUC 0.607 against 0.5, and the probabilities' range 0.1 to 0.343 against the threshold, and say what the model learned and what hid it. Then say what the confusion matrix at 0.25, [343 116 70 46], is in Lesson 3.4.3's terms, and what the runs at 0.646 to 0.673 say beside logistic regression's 0.665.

Hint: Three numbers say the model learned something and one threshold hides it. For the last part, ask what the hidden layer bought on the price and what it bought here.

Q4: The Loss That Fell to 0.70

Read this script and the output under it without running it, then say what the teammate who wrote it got right and wrong. nn.Sigmoid() is torch.sigmoid as a layer, so this model's output is a probability.

import torch
from torch import nn, optim
from sklearn.metrics import roc_auc_score
from sold_prep import X_fit, y_fit, X_val, y_val

loss_fn = nn.BCEWithLogitsLoss()
torch.manual_seed(42)
model = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 1), nn.Sigmoid())
optimizer = optim.SGD(model.parameters(), lr=0.1)
for epoch in range(1, 301):
    loss = loss_fn(model(X_fit), y_fit)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    if epoch in (1, 100, 300):
        model.eval()
        with torch.no_grad():
            val_loss = loss_fn(model(X_val), y_val).item()
        model.train()
        print(epoch, f"{loss.item():.4f}", f"{val_loss:.4f}")
model.eval()
with torch.no_grad():
    proba = model(X_val).squeeze(1)
print(round(roc_auc_score(y_val.squeeze(1).numpy(), proba.numpy()), 3))
print(round(proba.min().item(), 3), round(proba.max().item(), 3), int((proba >= 0.5).sum()))

1 0.8892 0.8903
100 0.7362 0.7368
300 0.7005 0.7008
0.474
0.008 0.059 0

The teammate reports that the loss fell from 0.89 to 0.70 and was still falling, so the model learns and only needs more epochs, that it calls nothing sold at 0.5 exactly as the Theory's model did, so the label's imbalance is the reason, and that the AUC must be a scoring mistake. Say what the loss saw at every epoch, why 0.70 is the number it fell to, what the outputs 0.008 to 0.059 are, what the AUC says, and the one-line fix.

Hint: Write out what reaches the loss for a row whose label is 0: a number between 0 and 1, then the loss's own sigmoid on it. Ask how low that row's charge can go, and compare with 0.6931.

Assignment

Assignment: The Weighted Loss, and What It Moves

The alert team reads the trained network's confusion matrix at 0.5, 459 true negatives, 116 false negatives and nothing else, and asks for a model whose own 0.5 calls sales, the way Lesson 3.4.3's class_weight="balanced" made logistic regression's predict call 261 listings instead of 3. nn.BCEWithLogitsLoss takes that lever as pos_weight, a tensor with one value per output column, which this lesson did not use: it multiplies the charge of every row whose label is 1, so the unsold count over the sold count makes the 360 sold fit rows weigh as much as the 1,365 unsold ones, the balance "balanced" struck. Write sold_weighted.py.

Print the two counts and their ratio to two decimals. On the seeded network before any step, print the weighted loss on the validation rows beside the plain one, then the weighted loss by hand, the formula with pos_weight multiplying the y * log(p) term, and the total weight the sold fit rows carry beside the unsold count. Write one train function that takes a loss object, trains the seeded network at 0.1 for the epochs asked, and at the reported epochs prints the training loss, the validation loss under the loss it trains with, the validation loss under the plain loss, and the AUC. Train with the weighted loss for 300 epochs, reporting at 1, 100 and 300, and print the probabilities' smallest, largest and mean beside the validation share sold, then at 0.5 the listings called, the sales among them, the accuracy and the four counts.

Train the plain loss for 300 epochs with the same function, find the plain model's threshold that calls exactly as many listings as the weighted model's 0.5 does, and print it with the listings it calls, the sales it finds and how many of its listings the weighted model also calls. Print both models' AUCs. Cut the weighted probabilities into the bands (0, 0.4], (0.4, 0.5], (0.5, 0.6] and (0.6, 1] with pd.cut and print each band's count, mean probability and share sold, Lesson 3.4.3's table. Finish with both losses trained for 1,000 epochs, reporting the last epoch, and say in a comment what the weights changed for certain and what they did not.

Hints:

• nn.BCEWithLogitsLoss(pos_weight=torch.tensor([unsold / sold])) is the whole lever, and (ratio * y_fit).sum() is the sold rows' total weight. The plain validation loss of a weighted model is the plain loss object called on the same logits.
• The plain model's matching threshold is the k-th largest of its probabilities, np.sort(plain)[::-1][k - 1], Lesson 3.4.3's cut, and (same & called).sum() counts the listings both models call.

import numpy as np
import pandas as pd
import torch
from torch import nn, optim
from sklearn.metrics import accuracy_score, confusion_matrix, roc_auc_score
from sold_prep import X_fit, y_fit, X_val, y_val

sold_val = y_val.squeeze(1).numpy()
sold = int(y_fit.sum().item())
unsold = len(y_fit) - sold
ratio = torch.tensor([unsold / sold])
print(unsold, sold, round(ratio.item(), 2))
plain_fn = nn.BCEWithLogitsLoss()
weighted_fn = nn.BCEWithLogitsLoss(pos_weight=ratio)

torch.manual_seed(42)
model = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 1))
with torch.no_grad():
    logits = model(X_val)
p = torch.sigmoid(logits)
print(f"{weighted_fn(logits, y_val).item():.4f} {plain_fn(logits, y_val).item():.4f}")
hand = -(ratio * y_val * torch.log(p) + (1 - y_val) * torch.log(1 - p)).mean()
print(f"{hand.item():.4f}")
print(round((ratio * y_fit).sum().item(), 1), unsold)

def train(loss_fn, epochs=300, report=()):
    torch.manual_seed(42)
    model = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 1))
    optimizer = optim.SGD(model.parameters(), lr=0.1)
    for epoch in range(1, epochs + 1):
        loss = loss_fn(model(X_fit), y_fit)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        model.eval()
        with torch.no_grad():
            out = model(X_val)
            proba = torch.sigmoid(out).squeeze(1).numpy()
        model.train()
        if epoch in report:
            scored = f"{loss_fn(out, y_val).item():.4f} {plain_fn(out, y_val).item():.4f}"
            print(epoch, f"{loss.item():.4f}", scored, round(roc_auc_score(sold_val, proba), 3))
    return proba

weighted = train(weighted_fn, report=(1, 100, 300))
print(round(weighted.min(), 3), round(weighted.max(), 3), round(weighted.mean(), 3), round(sold_val.mean(), 3))
called = weighted >= 0.5
found = int((called & (sold_val == 1)).sum())
print(int(called.sum()), found, round(accuracy_score(sold_val, called), 3), confusion_matrix(sold_val, called).ravel())

1365 360 3.79
1.0928 0.7530
1.0928
1365.0 1365
1 1.1054 1.0906 0.7476 0.527
100 1.0673 1.0565 0.6857 0.612
300 1.0406 1.0440 0.6753 0.636
0.26 0.725 0.493 0.202
248 69 0.607 [280 179  47  69]

plain = train(plain_fn)
cut = np.sort(plain)[::-1][called.sum() - 1]
same = plain >= cut
print(round(cut, 3), int(same.sum()), int((same & (sold_val == 1)).sum()), int((same & called).sum()))
print(round(roc_auc_score(sold_val, plain), 3), round(roc_auc_score(sold_val, weighted), 3))
bands = pd.cut(weighted, [0, 0.4, 0.5, 0.6, 1])
table = pd.DataFrame({"proba": weighted, "sold": sold_val})
print(table.groupby(bands, observed=True).size().tolist())
print(table.groupby(bands, observed=True).mean().round(3))

longer = train(weighted_fn, epochs=1000, report=(1000,))
longer_plain = train(plain_fn, epochs=1000, report=(1000,))
# The weights moved the probabilities for certain: a mean of 0.493 where 0.202
# sold, and 248 calls at 0.5 that a lower cut on the plain model nearly matches.
# They did not make the probabilities chances, and the AUC gap is these runs'.

0.216 248 66 219
0.607 0.636
[105, 222, 152, 96]
            proba   sold
(0.0, 0.4]  0.369  0.057
(0.4, 0.5]  0.447  0.185
(0.5, 0.6]  0.545  0.263
(0.6, 1.0]  0.650  0.302
1000 0.9539 1.0038 0.6440 0.673
1000 0.4777 0.4816 0.4816 0.646

Additional Resources

Optional, and nothing later depends on either. Worth opening if you want the full contract of the ten-class loss before Section 4.5, or the softmax as a layer with its one argument.

• Official docs - PyTorch 2.14, torch.nn.CrossEntropyLoss. The input as unnormalized logits of shape (n, C) that need not sum to 1, the target as class indices or, in a second form this lesson did not use, as a (n, C) tensor of class probabilities, and three arguments this lesson did not use: weight, one multiplier per class for an unbalanced training set, ignore_index, and label_smoothing.
• Official docs - PyTorch 2.14, torch.nn.Softmax. torch.softmax as a module: the formula, the dim argument along which every slice sums to 1, the shape rule that the output has the input's shape with every value in [0, 1], and the note that it is not the module to put in front of NLLLoss, a loss that takes log-probabilities, where LogSoftmax goes instead.

Summary

The classifier was the regression network with three things swapped, a label for the price, one loss object for another and a sigmoid read at the end, and the five steps never changed. What changed was what the numbers mean: a classifier's output is a logit, a number on a scale whose 0 is a coin, which becomes a probability only through a function the loss carries inside itself. Every wrong line in this lesson that ran without a word came from mistaking one of those scales for another, a probability handed to a loss that applies its own sigmoid, a softmax before a loss that applies its own, a threshold of 0.5 set above every probability the trained network returned on a label that sells one listing in five.

The honest number is the one the lesson ends on. On the price the hidden layer bought 7,140 of MAE, and on this label about a hundredth of AUC: 0.607 at the section's rate, 0.646 to 0.673 with more epochs or a faster rate, against 0.665 from the one-neuron model LogisticRegression fitted on the same columns, which the neuron trained by gradient steps climbed toward on its own. A network can bend where a line cannot, and whether that is worth anything on a label is found by running both and reporting both.

The form with one logit per class, a softmax and a target of class indices is the one Section 4.5 trains on 50,000 images, because ten classes do not fit in one logit, and Lesson 4.3.3 first breaks the full batch this loop still takes in one step into mini-batches, which those images need.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

import torch
from torch import nn, optim

y = torch.tensor(sold.to_numpy(), dtype=torch.float32).unsqueeze(1)
loss_fn = nn.BCEWithLogitsLoss()         # takes logits, float (n, 1) target
model = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 1))
optimizer = optim.SGD(model.parameters(), lr=0.1)
for epoch in range(1, 301):
    loss = loss_fn(model(X_fit), y_fit)  # logits in, never sigmoid(...)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
model.eval()
with torch.no_grad():
    proba = torch.sigmoid(model(X_val))  # probabilities, read here only
called = proba >= 0.25                   # the threshold is a choice

classes = y.squeeze(1).long()        # int64, (n,), 0 or 1
pair = nn.Sequential(nn.Linear(19, 32), nn.ReLU(), nn.Linear(32, 2))
ce = nn.CrossEntropyLoss()           # (n, 2) logits, (n,) classes
loss = ce(pair(X_fit), classes_fit)  # never ce(softmax(...), ...)
probs = torch.softmax(pair(X_val), dim=1)
pred = probs.argmax(dim=1)           # the class index per row