Тема 6 из 6

Данные и распознавание цифр

Последняя часть — данные. Настоящие наборы не помещаются в один тензор и требуют подготовки, поэтому в PyTorch для них есть Dataset и DataLoader.

Dataset

from torch.utils.data import Dataset


class PixelDigits(Dataset):
    def __init__(self, items):
        self.items = items              # например, список (путь к картинке, метка)

    def __len__(self):
        return len(self.items)          # сколько примеров

    def __getitem__(self, i):
        text, label = self.items[i]     # i-й пример: прочитать, подготовить
        x = torch.tensor([1.0 if ch == "#" else 0.0 for ch in text])
        return x, label

Dataset отвечает на два вопроса: «сколько примеров?» и «дай пример номер i». Картинки можно читать с диска прямо в __getitem__ — тогда в памяти не нужно держать весь набор. Для готовых тензоров есть TensorDataset(x, y).

DataLoader

from torch.utils.data import DataLoader

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=256, shuffle=False)

for xb, yb in train_loader:        # пакеты: xb — (64, …), yb — (64,)
    ...
  • batch_size — сколько примеров в пакете; последний пакет может быть меньше (drop_last=True его выбросит);
  • shuffle=True перемешивает примеры перед каждым проходом — для обучения; проверочные данные не перемешивают;
  • для повторяемости в DataLoader передают generator=torch.Generator().manual_seed(0).

Обучение по пакетам

for epoch in range(epochs):            # эпоха — один проход по всем данным
    model.train()
    for xb, yb in train_loader:
        optimizer.zero_grad()
        loss_fn(model(xb), yb).backward()
        optimizer.step()

Шаг по пакету, а не по всем данным сразу, — быстрее, экономит память и даже помогает сети обобщать.

Вероятности и уверенность

probs = torch.softmax(logits, dim=1)     # по классам: в каждой строке сумма 1
conf, digit = probs.max(dim=1)           # наибольшая вероятность и её класс

Логиты — не вероятности: их нельзя сравнивать с порогом вроде 0.5. После softmax каждая строка — распределение вероятностей по классам, и наибольшая из них — уверенность сети. Если она мала, честнее ответить «не знаю».

Что дальше

  • MNIST — 70 000 настоящих рукописных цифр 28×28: torchvision.datasets.MNIST скачает их сам;
  • свёрточные сети (nn.Conv2d) — для картинок они работают гораздо лучше полносвязных;
  • видеокарта: model.to("cuda") и xb.to("cuda") — если она есть;
  • сохранение: torch.save(model.state_dict(), "model.pt") и model.load_state_dict(torch.load("model.pt")).

Итог

  • Dataset: __len__ и __getitem__; TensorDataset;
  • DataLoader(batch_size, shuffle, drop_last); обучение по пакетам и эпохам;
  • softmax(dim=1) — вероятности, max(dim=1) — уверенность и класс;
  • eval() + no_grad() для предсказаний.