An Honest Evaluation Loop

Write the evaluation loop that runs at the end of every epoch.

Given a classifier and a DataLoader that yields (x, y) batches, return the accuracy, the macro F1 score and the confusion matrix over the whole dataset. Three details make this easy to get wrong:

1. Eval mode. Dropout and batch norm behave differently in training. Call model.eval() before the loop. Afterwards, put the model back in the mode it was in, because the training loop that called you expects that.
2. No graph. Wrap the loop in torch.no_grad(). Without it, every forward pass records an autograd graph that you will never use.
3. Accumulate counts, not scores. The last batch can be smaller, and F1 is not an average of per-batch F1 scores. Add up one confusion matrix over all batches, and compute every metric from it at the end.

The confusion matrix C has shape (n_classes, n_classes) and dtype int64. C[i, j] counts the examples whose true class is i and whose predicted class is j. For each class:

tp = C[c, c];  fp = C[:, c].sum() - tp;  fn = C[c, :].sum() - tp
F1_c = 2 * tp / (2 * tp + fp + fn)     # 0 when the denominator is 0
macro_f1 = mean over all n_classes of F1_c

Function signature

def evaluate(model, loader, n_classes):

Example

Input

# true labels      [0, 0, 1, 2, 2]
# predicted labels [0, 1, 1, 2, 0]   (argmax of the model's logits)
evaluate(model, loader, n_classes=3)

Output

{"accuracy": 0.6,
 "macro_f1": 0.6111,
 "confusion": tensor([[1, 1, 0],
                      [0, 1, 0],
                      [1, 0, 1]])}

Reasoning

Three of five predictions are right, so the accuracy is 0.6. Class 0 has tp = 1, fp = 1, fn = 1, so its F1 is 2 / 4 = 0.5. Class 1 has tp = 1, fp = 1, fn = 0, so its F1 is 2 / 3. Class 2 has tp = 1, fp = 0, fn = 1, so its F1 is 2 / 3. Their mean is 0.6111.

Constraints and notes

• The model returns logits of shape (batch, n_classes). The prediction is their argmax.
• accuracy and macro_f1 are Python floats. confusion is a tensor.
• A class that is never in the labels and never predicted has an F1 of 0. It still counts in the macro mean.