Introduction

In PyTorch image classification, model output is usually a numeric class index, not a human readable label. The missing piece is a stable mapping between class ids and label names during both training and inference. If that mapping is not saved and reused, predictions can look correct numerically but wrong semantically.

Build a Stable Label Mapping During Training

When using torchvision.datasets.ImageFolder, class names are derived from folder names and sorted alphabetically. This creates a deterministic class_to_idx dictionary that should be saved with the model checkpoint.

1import json
2from pathlib import Path
3import torch
4from torch import nn, optim
5from torch.utils.data import DataLoader
6from torchvision import datasets, models, transforms
7
8data_dir = Path("data/train")
9transform = transforms.Compose([
10    transforms.Resize((224, 224)),
11    transforms.ToTensor(),
12])
13
14dataset = datasets.ImageFolder(data_dir, transform=transform)
15loader = DataLoader(dataset, batch_size=32, shuffle=True)
16
17model = models.resnet18(weights=None)
18model.fc = nn.Linear(model.fc.in_features, len(dataset.classes))
19
20criterion = nn.CrossEntropyLoss()
21optimizer = optim.Adam(model.parameters(), lr=1e-3)
22
23model.train()
24for images, targets in loader:
25    optimizer.zero_grad()
26    logits = model(images)
27    loss = criterion(logits, targets)
28    loss.backward()
29    optimizer.step()
30
31checkpoint = {
32    "state_dict": model.state_dict(),
33    "class_to_idx": dataset.class_to_idx,
34}
35torch.save(checkpoint, "classifier.pt")
36
37with open("labels.json", "w", encoding="utf-8") as f:
38    json.dump(dataset.class_to_idx, f, indent=2)

Saving both model weights and label mapping ensures inference code can decode predictions consistently.

Decode Predicted Label at Inference Time

At inference, invert class_to_idx so predicted index maps back to label text. Then report probabilities for debugging and threshold tuning.

1import torch
2from torchvision import models, transforms
3from PIL import Image
4
5checkpoint = torch.load("classifier.pt", map_location="cpu")
6class_to_idx = checkpoint["class_to_idx"]
7idx_to_class = {v: k for k, v in class_to_idx.items()}
8
9model = models.resnet18(weights=None)
10model.fc = torch.nn.Linear(model.fc.in_features, len(class_to_idx))
11model.load_state_dict(checkpoint["state_dict"])
12model.eval()
13
14transform = transforms.Compose([
15    transforms.Resize((224, 224)),
16    transforms.ToTensor(),
17])
18
19image = Image.open("sample.jpg").convert("RGB")
20input_tensor = transform(image).unsqueeze(0)
21
22with torch.no_grad():
23    logits = model(input_tensor)
24    probs = torch.softmax(logits, dim=1)
25    top_prob, top_idx = torch.max(probs, dim=1)
26
27label = idx_to_class[int(top_idx.item())]
28confidence = float(top_prob.item())
29print(f"prediction: {label}, confidence: {confidence:.4f}")

If labels look swapped, the first thing to verify is that inference used the same mapping generated during training.

Multi Label Versus Single Label Clarification

CrossEntropyLoss with one class index per image solves single label classification. For multi label tasks, each image can have several labels and needs a different output and loss setup.

Use this rule:

• Single label: final layer size equals class count, target is one index, loss is CrossEntropyLoss
• Multi label: final layer size equals label count, target is multi hot vector, loss is BCEWithLogitsLoss

Choosing the wrong formulation leads to confusing label outputs even when training appears stable.

Persist Metadata for Deployment

Model files alone are not enough for production. Store preprocessing settings, image size, normalization values, and label map version in one metadata object.

1metadata = {
2    "image_size": [224, 224],
3    "normalization": "none",
4    "labels_version": "2026-03-04",
5    "class_to_idx": class_to_idx,
6}
7
8torch.save({"state_dict": model.state_dict(), "metadata": metadata}, "classifier_v2.pt")

This prevents silent drift between training and serving environments.

Common Pitfalls

A common mistake is rebuilding labels manually in inference code. Even one ordering difference can map class index zero to the wrong class name.

Another issue is applying different preprocessing between training and prediction. If resize and normalization differ, confidence and label quality can degrade sharply.

A third issue is treating a multi label dataset as single label and forcing one class output. The model then predicts the most dominant class and misses secondary labels.

Finally, avoid overwriting checkpoints without versioned metadata. Label map drift is hard to debug once historical artifacts are lost.

Summary

• Save class_to_idx during training and reuse it for every inference path
• Invert mapping at prediction time to decode class index into readable label
• Match loss function and output layer to single label or multi label task type
• Keep preprocessing identical between training and inference
• Version model metadata so label mappings stay reproducible in production