PyTorch One-Hot Encoding
One-hot encoding turns a class number into a row of zeros with a single 1. The 1 sits at the position of the class number.
Class numbers become rows
Three class numbers become three rows. Each row marks its class with a single 1.
labels = torch.tensor([0, 2, 1])
rows = F.one_hot(labels)
print(rows)num_classes fixes the width
Without num_classes, the width stops at the largest class present. The argument sets the width directly.
labels = torch.tensor([0, 2])
print("Default shape:", F.one_hot(labels).shape)
print("With num_classes=5:", F.one_hot(labels, num_classes=5).shape)Example
import torch
import torch.nn.functional as F
names = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
days = torch.arange(7)
rows = F.one_hot(days, num_classes=7)
print("Days:", names)
print(rows)
print("Row sums:", rows.sum(dim=1).tolist())Key points
- Use F.one_hot(labels) to turn class numbers into one-hot rows
- F: the short name for torch.nn.functional, a home for plain functions
- num_classes: the number of columns the rows should have
- Each one-hot row sums to 1
- One-hot rows suit inputs that give each class its own column
- Cross-entropy targets stay class numbers, not one-hot rows
In this example, the seven day names sit in a plain Python list. torch.arange(7) makes the day numbers 0 to 6. F.one_hot turns them into rows with num_classes=7. Each row marks one day with a single 1. The 1 moves one column right on each row. The row sums print as seven 1 values. Every day got exactly one mark. A model can now take the day as seven input columns.
Note
one_hot returns int64 rows, so cast with .float() before a layer uses them
Try it Yourself
Edit the code, then press Run.
Running is simulated ยท Python Playground
Exercise
What does F.one_hot(torch.tensor([2]), num_classes=4) return?