🔥 PyTorch

Полный справочник: тензоры и autograd, nn.Module, слои, Dataset и DataLoader, цикл обучения, оптимизаторы и планировщики, GPU и AMP, сохранение, torch.compile, экспорт, отладка и типичные ошибки.

Шпаргалки · ИИ · #pytorch #torch #deep-learning #python #gpu

Что это

PyTorch (Meta, Linux Foundation) фреймворк глубокого обучения: тензоры (как NumPy, но на GPU), автоматическое дифференцирование (autograd) и модули нейросетей (torch.nn). Динамический граф: код выполняется как обычный Python, удобно отлаживать.

pip install torch torchvision torchaudio            # CPU; для CUDA команда выбирается на pytorch.org/get-started
python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.backends.mps.is_available())"

import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, TensorDataset, random_split

Тензоры

x = torch.tensor([[1., 2.], [3., 4.]])                         # из данных (копирует)
torch.zeros(2, 3); torch.ones(2, 3); torch.full((2, 2), 7); torch.eye(3); torch.arange(10); torch.linspace(0, 1, 5)
torch.rand(2, 3); torch.randn(2, 3); torch.randint(0, 10, (2, 3)); torch.zeros_like(x); torch.empty(2, 2)
torch.manual_seed(42)                                           # воспроизводимость

x.shape; x.size(); x.ndim; x.dtype; x.device; x.numel(); x.requires_grad
x.float(); x.long(); x.half(); x.to(torch.float32); x.int()                              # типы
x.reshape(4); x.view(-1); x.flatten(); x.T; x.permute(1, 0); x.transpose(0, 1); x.unsqueeze(0); x.squeeze(); x.expand(3, 2, 2)
torch.cat([a, b], dim=0); torch.stack([a, b]); torch.split(x, 2); torch.chunk(x, 2)
x[0]; x[:, 1]; x[x > 2]; x[[0, 1]]; torch.where(x > 2, x, 0.); x.masked_fill(mask, 0); torch.gather(x, 1, idx)
a + b; a * b; a @ b; torch.matmul(a, b); a ** 2; x.sum(); x.mean(dim=0); x.max(dim=1); x.argmax(dim=1); x.softmax(dim=-1); torch.topk(x, 3)
x.numpy(); torch.from_numpy(arr); x.item(); x.tolist(); x.clone(); x.detach()

Типы: float32 (по умолчанию), float16 / bfloat16 (экономия памяти), int64 (индексы и метки классов), bool. Операции с суффиксом _ изменяют на месте (x.add_(1)). view требует непрерывной памяти, reshape справляется сам.

Устройства (CPU, CUDA, MPS)

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
x = x.to(device); model = model.to(device)                      # и данные, и модель на одном устройстве
torch.cuda.device_count(); torch.cuda.get_device_name(0); torch.cuda.memory_allocated() / 1e9; torch.cuda.empty_cache()

Копирование CPU↔GPU медленное: переносите батчи, а не отдельные числа; pin_memory=True и non_blocking=True ускоряют.

Autograd

w = torch.tensor(2.0, requires_grad=True)
loss = (w * 3 - 6) ** 2
loss.backward()                  # вычисляет градиенты
w.grad                           # d loss / d w
w.grad.zero_()                   # градиенты накапливаются: обнуляйте перед следующим шагом
with torch.no_grad(): y = model(x)            # без графа: инференс, экономия памяти
y = model(x).detach()                          # отсоединить от графа
torch.inference_mode()                         # ещё быстрее для инференса

Модель nn.Module

class Net(nn.Module):
    def __init__(self, d_in=10, d_hidden=64, n_classes=3, p=0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden), nn.BatchNorm1d(d_hidden), nn.ReLU(), nn.Dropout(p),
            nn.Linear(d_hidden, n_classes),
        )
    def forward(self, x):                 # вызывайте model(x), не model.forward(x)
        return self.net(x)

model = Net().to(device)
sum(p.numel() for p in model.parameters() if p.requires_grad)            # число параметров
model.state_dict(); model.parameters(); model.named_parameters(); model.train(); model.eval()

Выход классификатора это логиты (без softmax): CrossEntropyLoss применяет log_softmax сам.

Слои

Списки слоёв кладите в nn.ModuleList, а не в обычный list: иначе параметры не зарегистрируются.

class CNN(nn.Module):
    def __init__(self, n=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1),
        )
        self.head = nn.Linear(64, n)
    def forward(self, x): return self.head(self.features(x).flatten(1))

Форма изображений: (batch, channels, height, width).

Данные

class MyDataset(Dataset):
    def __init__(self, X, y, transform=None): self.X, self.y, self.transform = X, y, transform
    def __len__(self): return len(self.X)
    def __getitem__(self, i):
        x = self.X[i]
        if self.transform: x = self.transform(x)
        return x, self.y[i]

ds = TensorDataset(torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.long))
train_ds, val_ds = random_split(ds, [0.8, 0.2], generator=torch.Generator().manual_seed(42))
train_dl = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, persistent_workers=True)
val_dl = DataLoader(val_ds, batch_size=256)

from torchvision import datasets, transforms
tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
mnist = datasets.MNIST("data", train=True, download=True, transform=tf)

Аугментации: transforms.RandomHorizontalFlip, RandomCrop, ColorJitter, v2 API; тексты: torchtext, библиотеки Hugging Face datasets.

Цикл обучения

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

for epoch in range(epochs):
    model.train()
    for xb, yb in train_dl:
        xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(xb), yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)       # против взрыва градиентов
        optimizer.step()
    scheduler.step()

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for xb, yb in val_dl:
            pred = model(xb.to(device)).argmax(1).cpu()
            correct += (pred == yb).sum().item(); total += len(yb)
    print(f"epoch {epoch}: val_acc={correct / total:.3f}")

Порядок шага: обнулить градиенты → прямой проход → потеря → backward → step. model.train() включает Dropout и обновление BatchNorm, model.eval() выключает.

Потери и оптимизаторы

Оптимизаторы: SGD(momentum=0.9), Adam, AdamW (по умолчанию для большинства задач), RMSprop, Adagrad, LBFGS. Планировщики: StepLR, MultiStepLR, ExponentialLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau, LinearLR + SequentialLR (прогрев).

Ускорение: AMP, compile, градиенты

scaler = torch.amp.GradScaler("cuda")
with torch.autocast(device_type="cuda", dtype=torch.float16):          # bfloat16 не требует GradScaler
    loss = criterion(model(xb), yb)
scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()

model = torch.compile(model)                                            # PyTorch 2: ускорение через JIT-компиляцию
torch.backends.cudnn.benchmark = True                                   # для постоянных размеров входа
accum = 4                                                               # накопление градиентов: большой «виртуальный» батч
(loss / accum).backward(); (step % accum == 0) and (optimizer.step(), optimizer.zero_grad())

Дальше: DistributedDataParallel и torchrun (несколько GPU), FSDP (большие модели), torch.utils.checkpoint (экономия памяти), Hugging Face Accelerate, PyTorch Lightning.

Сохранение и загрузка

torch.save(model.state_dict(), "model.pt")                                      # рекомендуется: только веса
model = Net(); model.load_state_dict(torch.load("model.pt", map_location=device, weights_only=True)); model.eval()
torch.save({"epoch": e, "model": model.state_dict(), "opt": optimizer.state_dict()}, "ckpt.pt")      # чекпоинт для продолжения
scripted = torch.jit.script(model)                                              # TorchScript
torch.onnx.export(model, example_input, "model.onnx", opset_version=17)          # ONNX: перенос в другие рантаймы

Файлы .pt используют pickle: не загружайте чужие без weights_only=True. Форматы без pickle: safetensors.

Готовые модели и дообучение

from torchvision import models
m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
for p in m.parameters(): p.requires_grad = False                         # заморозить основу
m.fc = nn.Linear(m.fc.in_features, 5)                                    # новая голова
optimizer = torch.optim.AdamW(m.fc.parameters(), lr=1e-3)                # обучается только голова

Hugging Face Transformers: AutoModel.from_pretrained("bert-base-multilingual-cased"), PEFT / LoRA для экономного дообучения LLM, timm (модели зрения).

Отладка

torch.autograd.set_detect_anomaly(True)                                   # найти NaN в градиентах (медленно)
for n, p in model.named_parameters(): print(n, p.grad.norm() if p.grad is not None else None)
print(model); from torchinfo import summary; summary(model, input_size=(1, 10))
# проверка кода: обучите модель на одном маленьком батче, потеря должна упасть почти до нуля

Профилирование: torch.profiler, TensorBoard (torch.utils.tensorboard.SummaryWriter), Weights & Biases, MLflow.

Типичные ошибки

Воспроизводимость

torch.manual_seed(s), np.random.seed(s), random.seed(s), torch.use_deterministic_algorithms(True), DataLoader(worker_init_fn=...); полная детерминированность на GPU стоит скорости.

Экосистема

torchvision, torchaudio, Hugging Face (Transformers, Datasets, PEFT, Accelerate), timm, PyTorch Lightning, Ignite, torchmetrics, Captum (интерпретация), ONNX Runtime, TorchServe, vLLM (инференс LLM), DeepSpeed, Optuna, Weights & Biases.