๐ŸŽ“ A1Learner
๐Ÿ”

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?