PyTorch: Learn the Basics, часть 2 — Dataset, DataLoader и преобразования

Вторая часть перевода курса Learn the Basics (исходники — pytorch/tutorials, BSD-3-Clause). Здесь два раздела оригинала: «Datasets & DataLoaders» и «Transforms». Предыдущая часть — быстрый старт и тензоры, следующая — модель и autograd.

Как читать вместе с курсом ChinaAI

Свой класс Dataset — это обычный класс Python с «магическими» методами __len__ и __getitem__. Если эти слова пугают, освежите урок «Классы и магические методы». Таблица меток читается через pandas — его разбирает урок «pandas: таблицы как у профессионалов». В курсе ChinaAI Dataset и DataLoader впервые появляются в уроке «nn.Module, оптимизаторы, цикл обучения».

2. Dataset и DataLoader

Код обработки примеров данных легко превращается в запутанный и трудный для поддержки. В идеале код набора данных должен быть отделён от кода обучения модели — так его проще читать и переиспользовать. В PyTorch есть два базовых инструмента для работы с данными: torch.utils.data.DataLoader и torch.utils.data.Dataset. Они позволяют пользоваться и готовыми наборами данных, и вашими собственными. Dataset хранит примеры и соответствующие им метки, а загрузчик данныхзагрузчик данных (DataLoader)EN data loaderБерёт Dataset и выдаёт батчи: перемешивает, склеивает примеры в тензоры, может грузить данные в нескольких процессах. оборачивает Dataset в итерируемый объект, чтобы к примерам было удобно обращаться.

Предметные библиотеки PyTorch содержат множество готовых наборов данных (например, FashionMNIST). Они наследуют torch.utils.data.Dataset и реализуют функции, специфичные для конкретных данных. Ими удобно пользоваться, чтобы быстро опробовать модель и сравнить её с другими. Списки наборов: изображения, тексты и аудио.

Загрузка набора данных

Вот как загрузить набор Fashion-MNIST из TorchVision. Fashion-MNIST — это набор изображений товаров Zalando: 60 000 обучающих примеров и 10 000 тестовых. Каждый пример — чёрно-белое изображение 28×28 и метка одного из 10 классов.

Мы загружаем FashionMNIST Dataset со следующими параметрами:

• root — путь, где хранятся обучающие и тестовые данные;
• train — какую часть брать: обучающую или тестовую;
• download=True — скачать данные из интернета, если их нет в root;
• transform и target_transform — преобразования признаков и меток.

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import v2
import matplotlib.pyplot as plt


training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
)

Перебор и визуализация набора данных

К Dataset можно обращаться по индексу, как к списку: training_data[index]. С помощью matplotlib посмотрим на несколько примеров из обучающих данных.

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

(В оригинале здесь рисунок: сетка 3×3 из случайных картинок FashionMNIST — светлые силуэты одежды и обуви на чёрном фоне, над каждой подпись класса. У вас картинки будут другими: индексы выбираются случайно.)

---

Свой Dataset для собственных файлов

Свой класс Dataset обязан реализовать три функции: __init__, __len__ и __getitem__. Посмотрите на реализацию ниже: картинки FashionMNIST лежат в папке img_dir, а их метки — отдельно, в CSV-файле annotations_file.

В следующих подразделах разберём, что происходит в каждой из этих функций.

import os
import pandas as pd
from torchvision.io import decode_image

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = decode_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

__init__

Функция __init__ выполняется один раз, когда создаётся объект Dataset. В ней мы запоминаем папку с изображениями, файл с аннотациями и оба преобразования (о них подробнее — в следующем разделе).

Файл labels.csv выглядит так:

tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9

def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
    self.img_labels = pd.read_csv(annotations_file)
    self.img_dir = img_dir
    self.transform = transform
    self.target_transform = target_transform

__len__

Функция __len__ возвращает число примеров в наборе данных.

Пример:

def __len__(self):
    return len(self.img_labels)

__getitem__

Функция __getitem__ загружает и возвращает пример из набора данных по заданному индексу idx. По индексу она находит, где лежит изображение на диске, превращает его в тензор с помощью decode_image, берёт соответствующую метку из CSV-данных в self.img_labels, применяет к ним функции преобразования (если они заданы) и возвращает тензор-изображение и метку в виде кортежа.

def __getitem__(self, idx):
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    image = decode_image(img_path)
    label = self.img_labels.iloc[idx, 1]
    if self.transform:
        image = self.transform(image)
    if self.target_transform:
        label = self.target_transform(label)
    return image, label

Примечание переводчика: тот же приём для текста

В языковых моделях Dataset устроен так же, только вместо картинок — токены. Например, __len__ возвращает число окон длины block_size в длинном массиве токенов, а __getitem__(idx) — пару «окно токенов, то же окно со сдвигом на один токен». Именно так готовят данные в уроке «GPT на PyTorch».

---

Готовим данные к обучению с DataLoader

Dataset выдаёт признаки и метку по одному примеру за раз. При обучении модели обычно хочется подавать примеры мини-батчами, перемешивать данные в каждой эпохе, чтобы снизить переобучение, и ускорять загрузку с помощью модуля Python multiprocessing.

DataLoader — итерируемый объект, который прячет всю эту сложность за простым API.

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

Перебор DataLoader

Мы загрузили набор данных в DataLoader и можем перебирать его как нужно. Каждая итерация ниже возвращает батч train_features и train_labels (в них по batch_size=64 признаков и меток соответственно). Поскольку мы указали shuffle=True, после того как мы переберём все батчи, данные перемешаются (если нужен более тонкий контроль над порядком загрузки, посмотрите на Samplers).

# Показываем изображение и метку.
train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")
img = train_features[0].squeeze()
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"Label: {label}")

Вывод:

Feature batch shape: torch.Size([64, 1, 28, 28])
Labels batch shape: torch.Size([64])
Label: 1

(В оригинале под кодом — картинка первого примера из батча. Метка у вас может быть другой: батч перемешан. Метка 1 — это «Trouser», брюки.)

Примечание переводчика: num_workers

Ускорение загрузки в несколько процессов включается параметром num_workers, например DataLoader(training_data, batch_size=64, shuffle=True, num_workers=4). На маленьком FashionMNIST разницы почти не будет, а на больших наборах с декодированием картинок с диска — заметная. Если на Windows или в Jupyter с num_workers > 0 что-то падает, начните с num_workers=0.

Что почитать дальше

• API torch.utils.data

---

3. Преобразования (Transforms)

Данные не всегда приходят в том окончательном виде, который нужен для обучения алгоритмов машинного обучения. С помощью преобразований (transforms) мы изменяем данные и делаем их пригодными для обучения.

У всех наборов данных TorchVision есть два параметра: transform — чтобы изменять признаки, и target_transform — чтобы изменять метки. Оба принимают вызываемые объекты, в которых записана логика преобразования. Модуль torchvision.transforms предлагает несколько часто нужных преобразований из коробки.

Признаки FashionMNIST хранятся в формате PIL Image, а метки — целые числа. Для обучения нам нужны признаки в виде нормализованных тензоров, а метки — в виде тензоров в кодировке one-hot кодированиеone-hot кодирование (унитарный код)EN one-hot encodingВектор длины V из нулей с единицей в позиции номера токена. Все one-hot векторы попарно ортогональны.. Для этих преобразований мы используем API torchvision.transforms.v2 вместе с torch.nn.functional.one_hot.

import torch
import torch.nn.functional as F
from torchvision import datasets
from torchvision.transforms import v2

ds = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)]),
    target_transform=v2.Lambda(
        lambda y: F.one_hot(torch.tensor(y), num_classes=10).float()
    ),
)

ToImage() и ToDtype()

API torchvision.transforms.v2 заменяет устаревшее преобразование ToTensor конвейером из двух шагов. v2.ToImage превращает изображение PIL или массив NumPy ndarray в тензор torchvision.tv_tensors.Image, а v2.ToDtype с scale=True приводит его к типу float32 и масштабирует яркость пикселей в диапазон [0., 1.].

(Примечание переводчика: пиксель хранится как целое число от 0 до 255; после scale=True значение 255 станет 1.0, а 51 — 0.2. Нейросети учатся заметно устойчивее на маленьких числах.)

Lambda-преобразования

Lambda-преобразование применяет любую заданную пользователем lambda-функцию. Здесь мы с помощью torch.nn.functional.one_hot превращаем целочисленную метку в тензор длины 10 (по числу классов в наборе) в кодировке one-hot, а затем приводим его к типу float, чтобы тип совпадал с ожидаемым.

target_transform = v2.Lambda(
    lambda y: F.one_hot(torch.tensor(y), num_classes=10).float()
)

(Примечание переводчика: для метки 3 («Dress») получится tensor([0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]). Обратите внимание: в быстром старте и в разделе 6 метки остаются целыми числами, потому что nn.CrossEntropyLoss принимает номер класса напрямую. One-hot здесь — пример того, как работает target_transform.)

Что почитать дальше

• Getting started with transforms v2
• API torchvision.transforms.v2

---

Словарик

EN	RU
dataset	набор данных
sample, label	пример, метка
features	признаки
annotations file	файл аннотаций (разметки)
minibatch	мини-батч
shuffle	перемешивание
transform	преобразование
one-hot encoding	кодирование one-hot

Дальше: часть 3 — сборка модели и автоматическое дифференцирование.