PyTorch: Learn the Basics, часть 1 — быстрый старт и тензоры

Это перевод вводного курса Learn the Basics из официальных туториалов PyTorch. Исходный текст — файлы beginner_source/basics/*.py в репозитории pytorch/tutorials (лицензия BSD-3-Clause); выводы программ взяты с опубликованных страниц туториалов, где код выполнен на GPU. Курс разбит у нас на четыре страницы:

1. Быстрый старт и тензоры — эта страница.
2. Данные: Dataset, DataLoader и преобразования.
3. Модель и автоматическое дифференцирование.
4. Цикл обучения, сохранение и загрузка.

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

Курс рассчитан на тех, кто уже немного знает Python и идеи глубокого обучения. У нас это ступень 3: читайте его параллельно с уроками «Тензоры и autograd в PyTorch» и «nn.Module, оптимизаторы, цикл обучения». Раздел о тензорах опирается на NumPy — если массивы и срезы ещё непривычны, сначала пройдите «Массивы, формы, индексация». Пометки «Примечание переводчика» — наши, остальное — перевод оригинала.

Изучаем основы

Авторы: Сурадж Субраманиан, Сет Хуарес, Кэсси Бревиу, Дмитрий Сошников, Ари Борнштейн.

Большинство задач машинного обучения сводятся к одному и тому же: работать с данными, создавать модели, оптимизировать параметры модели и сохранять обученные модели. Этот туториал знакомит с полным циклом машинного обучения на PyTorch и даёт ссылки, чтобы подробнее изучить каждый шаг.

Мы обучим нейросеть на наборе данных FashionMNIST: она будет предсказывать, к какому из классов относится входное изображение: футболка/топ (T-shirt/top), брюки (Trouser), пуловер (Pullover), платье (Dress), пальто (Coat), сандалия (Sandal), рубашка (Shirt), кроссовок (Sneaker), сумка (Bag) или ботинок (Ankle boot).

Туториал предполагает базовое знакомство с Python и понятиями глубокого обучения.

Как запустить код туториала

Код можно запустить двумя способами:

• В облаке. Самый простой способ начать. В начале каждого раздела оригинала есть ссылка «Run in Google Colab»: она открывает готовый ноутбук в Google Colab, где всё уже установлено.
• Локально. Для этого нужно установить PyTorch и TorchVision на свой компьютер (инструкция по установке). Скачайте ноутбук или скопируйте код в свою среду разработки.

Примечание переводчика: локальный запуск

В виртуальном окружении курса (см. урок «Python, виртуальные окружения и Jupyter») PyTorch уже стоит, а пакета torchvision может не быть. Поставьте его командой pip install torchvision — он нужен для набора FashionMNIST и преобразований изображений. Сам набор (около 30 МБ) код скачает в папку data/ при первом запуске. Всё в этом туториале обучается на процессоре за несколько минут, видеокарта не обязательна.

Как пользоваться этим руководством

Если вы знакомы с другими фреймворками глубокого обучения, начните с раздела «0. Быстрый старт» — он быстро познакомит вас с API PyTorch.

Если фреймворки глубокого обучения для вас в новинку, переходите сразу к первому разделу пошагового руководства — «1. Тензоры».

Содержание:

0. Быстрый старт (ниже на этой странице)
1. Тензоры (ниже на этой странице)
2. Dataset и DataLoader
3. Преобразования (Transforms)
4. Сборка модели
5. Автоматическое дифференцирование
6. Цикл оптимизации
7. Сохранение, загрузка и использование модели

---

0. Быстрый старт

В этом разделе мы пробежимся по API для типичных задач машинного обучения. Чтобы разобраться глубже, переходите по ссылкам в каждом подразделе.

Работа с данными

У PyTorch есть два базовых инструмента для работы с данными: torch.utils.data.DataLoader и torch.utils.data.Dataset. Dataset хранит примеры и соответствующие им метки, а DataLoader оборачивает Dataset в итерируемый объект — его можно перебирать в цикле for.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import v2

У PyTorch есть предметные библиотеки — TorchText, TorchVision и TorchAudio, — и в каждой есть наборы данных. В этом туториале мы возьмём набор из TorchVision.

Модуль torchvision.datasets содержит объекты Dataset для многих реальных наборов данных компьютерного зрения, например CIFAR и COCO (полный список). Здесь мы используем FashionMNIST. Каждый Dataset из TorchVision принимает два аргумента — transform и target_transform, — которые изменяют примеры и метки соответственно.

# Скачиваем обучающие данные из открытых наборов.
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 как аргумент в DataLoader. Тот оборачивает набор данных в итерируемый объект и умеет автоматически собирать батчи, выбирать примеры, перемешивать их и загружать данные в несколько процессов. Здесь мы задаём размер батча 64: каждый элемент, который выдаёт DataLoader, — это батч из 64 признаков и 64 меток.

batch_size = 64

# Создаём загрузчики данных.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

Вывод:

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64

(Примечание переводчика: [N, C, H, W] — это «число примеров в батче, число каналов, высота, ширина». Картинки FashionMNIST чёрно-белые, поэтому канал один.)

Подробнее — в разделе о загрузке данных в PyTorch.

---

Создание моделей

Чтобы задать нейросеть в PyTorch, мы создаём класс, наследующий nn.Module. Слои сети мы определяем в функции __init__, а то, как данные проходят через сеть, — в функции forward. Чтобы ускорить вычисления, мы переносим сеть на ускоритель (accelerator) — CUDA, MPS, MTIA или XPU. Если ускоритель доступен, используем его, иначе — процессор (CPU).

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")

# Определяем модель
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().to(device)
print(model)

Вывод:

Using cuda device
NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

(Примечание переводчика: без видеокарты первая строка будет Using cpu device — это нормально. Функция torch.accelerator появилась в PyTorch 2.6; в более старых версиях пишут device = "cuda" if torch.cuda.is_available() else "cpu".)

Подробнее — в разделе о сборке нейросетей в PyTorch.

---

Оптимизация параметров модели

Чтобы обучить модель, нужны функция потерь (loss function) и оптимизатор (optimizer).

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

За один проход цикла обучения модель делает предсказания на обучающем наборе (данные подаются батчами) и с помощью обратного распространения ошибки (backpropagation) подправляет свои параметры.

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Считаем ошибку предсказания
        pred = model(X)
        loss = loss_fn(pred, y)

        # Обратное распространение ошибки
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

Ещё мы проверяем качество модели на тестовом наборе, чтобы убедиться, что она учится.

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

Обучение идёт в несколько итераций — эпохэпохаEN epochОдин полный проход по обучающим данным. LLM на претрейне часто не делают и одной эпохи. (epochs). За каждую эпоху модель подбирает параметры, чтобы предсказывать лучше. После каждой эпохи мы печатаем точность (accuracy) и потери модели; хочется видеть, что точность растёт, а потери падают с каждой эпохой.

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

Вывод (начало и итоги эпох):

Epoch 1
-------------------------------
loss: 2.305831  [   64/60000]
loss: 2.286976  [ 6464/60000]
loss: 2.269019  [12864/60000]
...
loss: 2.171268  [57664/60000]
Test Error:
 Accuracy: 53.1%, Avg loss: 2.158626

Epoch 2
...
 Accuracy: 61.5%, Avg loss: 1.887746
Epoch 3
...
 Accuracy: 62.6%, Avg loss: 1.508443
Epoch 4
...
 Accuracy: 63.6%, Avg loss: 1.234293
Epoch 5
...
 Accuracy: 65.0%, Avg loss: 1.070535

Done!

(Примечание переводчика: ваши числа будут немного другими — начальные веса случайны. 65% после пяти эпох — немного: SGD с шагом 1e-3 учится медленно. В разделе 6 то же обучение идёт 10 эпох и доходит до 71%. Попробуйте заменить оптимизатор на torch.optim.Adam или увеличить шаг обучения и сравните кривые — это хорошее первое упражнение.)

Подробнее — в разделе об обучении модели.

---

Сохранение моделей

Обычный способ сохранить модель — сериализовать её внутренний словарь состояния (state dictionary), в котором лежат параметры модели.

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

Вывод:

Saved PyTorch Model State to model.pth

Загрузка моделей

Чтобы загрузить модель, нужно заново создать её структуру и загрузить в неё словарь состояния.

model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))

Вывод:

<All keys matched successfully>

Теперь модель можно использовать для предсказаний.

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
    x = x.to(device)
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

Вывод:

Predicted: "Ankle boot", Actual: "Ankle boot"

Подробнее — в разделе о сохранении и загрузке модели.

---

1. Тензоры

ТензорытензорEN tensorВ ML — просто многомерный массив чисел. torch.Tensor устроен как np.ndarray, но умеет работать на GPU и считать градиенты. — это особая структура данных, очень похожая на массивы и матрицы. В PyTorch тензорами кодируют входы и выходы модели, а также её параметры.

Тензоры похожи на ndarray из NumPy, но тензоры умеют работать на GPU и других аппаратных ускорителях. Более того, тензор и массив NumPy часто могут использовать одну и ту же область памяти, и тогда копировать данные не нужно (см. раздел «Мост с NumPy»). Ещё тензоры оптимизированы для автоматического дифференцирования (подробнее — в разделе Autograd). Если вы знакомы с ndarray, API тензоров покажется родным. Если нет — читайте дальше!

import torch
import numpy as np

Создание тензора

Тензор можно создать разными способами. Посмотрите на примеры.

Прямо из данных

Тензор можно создать прямо из данных. Тип данных определяется автоматически.

data = [[1, 2],[3, 4]]
x_data = torch.tensor(data)

Из массива NumPy

Тензор можно создать из массива NumPy (и наоборот — см. «Мост с NumPy»).

np_array = np.array(data)
x_np = torch.from_numpy(np_array)

Из другого тензора

Новый тензор сохраняет свойства (форму, тип данных) тензора-аргумента, если их явно не переопределить.

x_ones = torch.ones_like(x_data) # сохраняет свойства x_data
print(f"Ones Tensor: \n {x_ones} \n")

x_rand = torch.rand_like(x_data, dtype=torch.float) # переопределяет тип данных x_data
print(f"Random Tensor: \n {x_rand} \n")

Вывод:

Ones Tensor:
 tensor([[1, 1],
        [1, 1]])

Random Tensor:
 tensor([[0.8226, 0.9592],
        [0.8721, 0.7620]])

Со случайными или постоянными значениями

shape — это кортеж с размерами тензора. В функциях ниже он задаёт размерность создаваемого тензора.

shape = (2,3)
rand_tensor = torch.rand(shape)
ones_tensor = torch.ones(shape)
zeros_tensor = torch.zeros(shape)

print(f"Random Tensor: \n {rand_tensor} \n")
print(f"Ones Tensor: \n {ones_tensor} \n")
print(f"Zeros Tensor: \n {zeros_tensor}")

Вывод:

Random Tensor:
 tensor([[0.1857, 0.4198, 0.1457],
        [0.6639, 0.2573, 0.4555]])

Ones Tensor:
 tensor([[1., 1., 1.],
        [1., 1., 1.]])

Zeros Tensor:
 tensor([[0., 0., 0.],
        [0., 0., 0.]])

---

Атрибуты тензора

Атрибуты тензора описывают его форму, тип данных и устройство, на котором он хранится.

tensor = torch.rand(3,4)

print(f"Shape of tensor: {tensor.shape}")
print(f"Datatype of tensor: {tensor.dtype}")
print(f"Device tensor is stored on: {tensor.device}")

Вывод:

Shape of tensor: torch.Size([3, 4])
Datatype of tensor: torch.float32
Device tensor is stored on: cpu

---

Операции над тензорами

Более 1200 операций над тензорами — арифметика, линейная алгебра, работа с матрицами (транспонирование, индексация, срезы), сэмплирование и многое другое — подробно описаны в документации.

Любую из этих операций можно выполнить на CPU и на ускорителе: CUDA, MPS, MTIA или XPU. Если вы работаете в Colab, подключите ускоритель через меню Runtime → Change runtime type → GPU.

По умолчанию тензоры создаются на CPU. На ускоритель их нужно переносить явно методом .to (предварительно проверив, что ускоритель доступен). Помните, что копирование больших тензоров между устройствами может дорого обходиться и по времени, и по памяти!

# Переносим тензор на текущий ускоритель, если он есть
if torch.accelerator.is_available():
    tensor = tensor.to(torch.accelerator.current_accelerator())

Попробуйте некоторые операции из списка. Если вы знакомы с API NumPy, работать с API тензоров будет легко.

Стандартная индексация и срезы, как в NumPy

tensor = torch.ones(4, 4)
print(f"First row: {tensor[0]}")
print(f"First column: {tensor[:, 0]}")
print(f"Last column: {tensor[..., -1]}")
tensor[:,1] = 0
print(tensor)

Вывод:

First row: tensor([1., 1., 1., 1.])
First column: tensor([1., 1., 1., 1.])
Last column: tensor([1., 1., 1., 1.])
tensor([[1., 0., 1., 1.],
        [1., 0., 1., 1.],
        [1., 0., 1., 1.],
        [1., 0., 1., 1.]])

Соединение тензоров. Функция torch.cat склеивает последовательность тензоров вдоль заданного измерения. Посмотрите также на torch.stack — другую функцию соединения, которая немного отличается от torch.cat.

t1 = torch.cat([tensor, tensor, tensor], dim=1)
print(t1)

Вывод:

tensor([[1., 0., 1., 1., 1., 0., 1., 1., 1., 0., 1., 1.],
        [1., 0., 1., 1., 1., 0., 1., 1., 1., 0., 1., 1.],
        [1., 0., 1., 1., 1., 0., 1., 1., 1., 0., 1., 1.],
        [1., 0., 1., 1., 1., 0., 1., 1., 1., 0., 1., 1.]])

(Примечание переводчика: cat склеивает вдоль существующего измерения — здесь три тензора 4×4 дали 4×12. stack добавляет новое измерение: torch.stack([tensor, tensor, tensor]) даст форму 3×4×4.)

Арифметические операции

# Матричное умножение двух тензоров. y1, y2, y3 получатся одинаковыми.
# ``tensor.T`` возвращает транспонированный тензор
y1 = tensor @ tensor.T
y2 = tensor.matmul(tensor.T)

y3 = torch.rand_like(y1)
torch.matmul(tensor, tensor.T, out=y3)


# Поэлементное произведение. z1, z2, z3 получатся одинаковыми.
z1 = tensor * tensor
z2 = tensor.mul(tensor)

z3 = torch.rand_like(tensor)
torch.mul(tensor, tensor, out=z3)

Вывод (результат последней строки, z3):

tensor([[1., 0., 1., 1.],
        [1., 0., 1., 1.],
        [1., 0., 1., 1.],
        [1., 0., 1., 1.]])

Тензоры из одного элемента. Если у вас тензор из одного элемента — например, после того как все значения тензора сложили в одно, — его можно превратить в обычное число Python методом item():

agg = tensor.sum()
agg_item = agg.item()
print(agg_item, type(agg_item))

Вывод:

12.0 <class 'float'>

Операции на месте (in-place). Операции, которые записывают результат в сам операнд, называются операциями на месте. Их имена оканчиваются на _. Например, x.copy_(y) и x.t_() изменят x.

print(f"{tensor} \n")
tensor.add_(5)
print(tensor)

Вывод:

tensor([[1., 0., 1., 1.],
        [1., 0., 1., 1.],
        [1., 0., 1., 1.],
        [1., 0., 1., 1.]])

tensor([[6., 5., 6., 6.],
        [6., 5., 6., 6.],
        [6., 5., 6., 6.],
        [6., 5., 6., 6.]])

Заметка

Операции на месте немного экономят память, но могут мешать при вычислении производных: история вычислений сразу теряется. Поэтому пользоваться ими не рекомендуется.

---

Мост с NumPy

Тензоры на CPU и массивы NumPy могут использовать одну и ту же область памяти, и тогда изменение одного меняет и другое.

Из тензора в массив NumPy

t = torch.ones(5)
print(f"t: {t}")
n = t.numpy()
print(f"n: {n}")

Вывод:

t: tensor([1., 1., 1., 1., 1.])
n: [1. 1. 1. 1. 1.]

Изменение тензора отражается в массиве NumPy.

t.add_(1)
print(f"t: {t}")
print(f"n: {n}")

Вывод:

t: tensor([2., 2., 2., 2., 2.])
n: [2. 2. 2. 2. 2.]

Из массива NumPy в тензор

n = np.ones(5)
t = torch.from_numpy(n)

Изменения в массиве NumPy отражаются в тензоре.

np.add(n, 1, out=n)
print(f"t: {t}")
print(f"n: {n}")

Вывод:

t: tensor([2., 2., 2., 2., 2.], dtype=torch.float64)
n: [2. 2. 2. 2. 2.]

(Примечание переводчика: обратите внимание на dtype=torch.float64. NumPy по умолчанию создаёт числа двойной точности, и from_numpy сохраняет этот тип. Нейросети обычно работают во float32, поэтому такие тензоры часто приводят методом .float(). Общая память есть только у тензоров на CPU: после .to("cuda") получится копия.)

---

Словарик

EN	RU
tensor	тензор
shape, dtype, device	форма, тип данных, устройство
accelerator	ускоритель (GPU и подобные)
batch	батч, мини-батч
epoch	эпоха
loss function	функция потерь
optimizer	оптимизатор
backpropagation	обратное распространение ошибки
state dictionary (state_dict)	словарь состояния
in-place operation	операция на месте

Дальше: часть 2 — Dataset, DataLoader и преобразования.