PyTorch: тензоры

Массивы NumPy отлично подходят для численных расчётов, но для глубокого обучения им не хватает двух вещей: они не работают на GPU и не умеют автоматически вычислять производные. PyTorch закрывает именно эти два пробела. Его главный объект — тензор — очень похож на массив NumPy, но может находиться на видеокарте и отслеживать свои градиенты. В этом уроке ты научишься свободно работать с тензорами, а следующий урок посвящён автоматическому дифференцированию.

Что такое PyTorch и как его установить

PyTorch — библиотека глубокого обучения с открытым исходным кодом, созданная в лаборатории искусственного интеллекта компании Meta (ранее Facebook); с 2022 года ею управляет PyTorch Foundation в составе Linux Foundation. Исследователи и компании широко используют её, чтобы создавать, обучать и внедрять нейронные сети. PyTorch не работает в браузере, поэтому код PyTorch из этого модуля запускай на своём компьютере; результат показан под каждым примером.

1. 1
   Создай виртуальное окружение

   В папке проекта выполни python -m venv .venv и активируй окружение — так PyTorch не будет мешать другим проектам.

2. 2
   Выбери сборку для CPU или CUDA

   Если видеокарты NVIDIA нет, сборки для CPU для этого модуля вполне достаточно. Для сборки с GPU (CUDA) возьми точную команду из конфигуратора на сайте pytorch.org, выбрав свою ОС и версию CUDA.

3. 3
   Установи и проверь

   Выполни pip install torch, затем в Python напиши import torch и выведи версию.

python -m venv .venv
# Windows: .venv\Scripts\activate    macOS/Linux: source .venv/bin/activate
pip install torch numpy
# CPU-only build (smaller download):
pip install torch --index-url https://download.pytorch.org/whl/cpu

import torch

print(torch.__version__)
print(torch.cuda.is_available())

2.14.0+cpu
False

Создание тензоров: shape и dtype

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(a)
print(a.shape, a.ndim, a.dtype)
print(torch.tensor([1.5, 2.0]).dtype)
print(torch.zeros(2, 3))
print(torch.arange(0, 10, 3))
torch.manual_seed(42)
print(torch.rand(2, 2))
print(a.float().dtype, a.to(torch.float16).dtype)

tensor([[1, 2, 3],
        [4, 5, 6]])
torch.Size([2, 3]) 2 torch.int64
torch.float32
tensor([[0., 0., 0.],
        [0., 0., 0.]])
tensor([0, 3, 6, 9])
tensor([[0.8823, 0.9150],
        [0.3829, 0.9593]])
torch.float32 torch.float16

• torch.tensor(...) сам определяет тип: целые → int64, дробные → float32. А в NumPy для дробных по умолчанию float64!
• torch.zeros, torch.ones, torch.arange, torch.linspace создают готовые тензоры, а torch.rand (равномерное распределение от 0 до 1) и torch.randn (стандартное нормальное) — случайные.
• torch.manual_seed(42) фиксирует генератор случайных чисел — результаты становятся воспроизводимыми.
• .float(), .long() или .to(torch.float16) меняют тип и возвращают новый тензор.

Изменение формы и индексация

import torch

m = torch.arange(12).reshape(3, 4)
print(m)
print(m[0], m[:, 1])
print(m[1, 2].item())
print(m[m > 8])
v = m.view(2, 6)
print(v.shape, m.T.shape)
print(m.unsqueeze(0).shape, m.flatten().shape)
v[0, 0] = 100
print(m[0, 0])

tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
tensor([0, 1, 2, 3]) tensor([1, 5, 9])
6
tensor([ 9, 10, 11])
torch.Size([2, 6]) torch.Size([4, 3])
torch.Size([1, 3, 4]) torch.Size([12])
tensor(100)

• Индексация как в NumPy: m[0] — первая строка, m[:, 1] — второй столбец, m[m > 8] — выбор по маске. .item() превращает тензор из одного элемента в обычное число Python.
• reshape(2, 6) и view(2, 6) меняют форму, не меняя числа элементов; если вместо одного размера написать -1, PyTorch вычислит его сам: x.view(-1, 784).
• view не копирует данные, а смотрит на ту же память через другое «окно». Поэтому v[0, 0] = 100 изменило и m!
• unsqueeze(0) добавляет новую ось размера 1 (например, чтобы превратить одно изображение в пакет), squeeze() удаляет такие оси, .T транспонирует матрицу, а permute переставляет оси в любом порядке.

Операции: транслирование и матричное умножение

Операторы +, -, *, /, ** работают поэлементно. Редукции sum, mean, max считаются вдоль оси, заданной аргументом dim: a.mean(dim=0) — среднее по каждому столбцу. Тензоры разной формы согласуются по правилам транслирования (broadcasting):

1. Формы сравниваются справа налево; более короткая форма дополняется слева осями размера 1.
2. Два размера совместимы, если они равны или один из них равен 1.
3. Ось размера 1 «растягивается» до размера другой (без копирования данных).

import torch

a = torch.tensor([[1., 2.], [3., 4.]])
b = torch.tensor([10., 20.])
print(a + b)
print(a * a)
print(a @ a)
print(a.sum(), a.mean(dim=0))
X = torch.ones(64, 3)
W = torch.ones(3, 5)
print((X @ W).shape)
print((torch.ones(4, 1) + torch.ones(3)).shape)

tensor([[11., 22.],
        [13., 24.]])
tensor([[ 1.,  4.],
        [ 9., 16.]])
tensor([[ 7., 10.],
        [15., 22.]])
tensor(10.) tensor([2., 3.])
torch.Size([64, 5])
torch.Size([4, 3])

Работа с GPU и NumPy

В видеокарте тысячи простых ядер, которые считают матричные произведения параллельно; обучение больших моделей на GPU часто в разы, а то и в десятки раз быстрее, чем на CPU. Стандартный приём: в начале скрипта выбрать device, а затем переносить туда модель и данные с помощью .to(device). На компьютерах Apple Silicon для GPU есть устройство 'mps'.

import numpy as np
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
x = torch.ones(2, 2, device=device)
print(device, x.device)
arr = np.array([1.0, 2.0, 3.0])
t = torch.from_numpy(arr)
arr[0] = 99
print(t)
back = (t * 2).numpy()
print(back, type(back).__name__)

cpu cpu
tensor([99.,  2.,  3.], dtype=torch.float64)
[198.   4.   6.] ndarray

torch.from_numpy(arr) разделяет память массива: изменение arr[0] = 99 видно и в тензоре. Обрати внимание, что dtype остался float64 — перед обучением переведи его во float32 через .float(). В обратную сторону .numpy() работает только для тензоров на CPU: тензор с GPU сначала верни через .cpu(), а если он отслеживает градиенты, сначала вызови .detach(): t.detach().cpu().numpy().