Основы PyTorch

🎯 Цель

После прочтения этой главы:

• Узнаете API PyTorch: tensor, autograd, nn.Module, DataLoader
• Сможете создавать свою модель через nn.Module
• Сможете писать полный training loop (на CPU или GPU)
• Узнаете, как сохранить, загрузить модель и сделать inference
• Познакомитесь с выводом в production (torch.jit, ONNX)

Что нужно изучить

• Tensor — NumPy ndarray + GPU support + autograd
• Autograd — автоматическая дифференциация
• nn.Module — построение модели
• nn.Linear, nn.Conv2d, nn.RNN — слои
• Loss functions — nn.MSELoss, nn.CrossEntropyLoss и т.д.
• Optimizers — optim.SGD, optim.Adam
• Dataset и DataLoader — batch loading
• Device management — CPU/GPU/MPS
• Сохранение/загрузка — state_dict
• TorchScript — production export

Библиотеки

# CPU
pip install torch torchvision torchaudio

# CUDA 12.1 (NVIDIA GPU)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# Mac M1/M2/M3 — default install работает с MPS

Проверка:

import torch
print(torch.__version__)
print(torch.cuda.is_available())     # NVIDIA GPU
print(torch.backends.mps.is_available())  # Mac

Важные темы

Tensor — сердце PyTorch

import torch

# Создание
a = torch.tensor([1, 2, 3])              # int64
b = torch.tensor([1.0, 2.0, 3.0])        # float32
c = torch.zeros(3, 4)                    # 3x4 нулей
d = torch.randn(2, 3)                    # normal random
e = torch.arange(10)                     # [0..9]

# Из NumPy
import numpy as np
arr = np.array([1, 2, 3])
t = torch.from_numpy(arr)                # share memory!
arr_back = t.numpy()                     # share memory!

# Атрибуты
print(a.shape, a.dtype, a.device)        # torch.Size([3]) torch.int64 cpu

Device management

# Автоматический выбор лучшего device
device = "cuda" if torch.cuda.is_available() else (
    "mps" if torch.backends.mps.is_available() else "cpu"
)

# Перенос tensor на device
x = torch.randn(1000, 1000).to(device)
model = MyModel().to(device)

# Внимание: оба tensor должны быть на одном device
# y = x @ y  # ❌ если y на CPU
# y = x @ y.to(device)  # ✅

Autograd — magic

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1     # y = x² + 3x + 1
y.backward()                # dy/dx = 2x + 3
print(x.grad)               # tensor(7.0)  ← при x=2: 2*2+3=7

# Основной механизм — строится computational graph и при вызове backward
# для каждого x вычисляется gradient

Паттерн nn.Module

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(0.3)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        x = self.activation(self.fc1(x))
        x = self.dropout(x)
        x = self.activation(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

model = MyModel(input_dim=784, hidden_dim=256, output_dim=10)
print(model)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

Dataset и DataLoader

from torch.utils.data import Dataset, DataLoader

class CSVDataset(Dataset):
    def __init__(self, csv_path):
        df = pd.read_csv(csv_path)
        self.X = torch.tensor(df.drop("target", axis=1).values, dtype=torch.float32)
        self.y = torch.tensor(df["target"].values, dtype=torch.long)
    
    def __len__(self):
        return len(self.y)
    
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

train_dataset = CSVDataset("train.csv")
train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,       # параллельная загрузка данных
    pin_memory=True,     # быстрее для GPU
)

Примеры кода

Полный training loop

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

device = "cuda" if torch.cuda.is_available() else "cpu"

# Модель
model = MyModel(784, 256, 10).to(device)

# Loss и optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Training loop
def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss, total_correct, total = 0, 0, 0
    
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        
        optimizer.zero_grad()
        logits = model(X)
        loss = criterion(logits, y)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item() * X.size(0)
        total_correct += (logits.argmax(dim=1) == y).sum().item()
        total += X.size(0)
    
    return total_loss / total, total_correct / total

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, total_correct, total = 0, 0, 0
    
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        logits = model(X)
        loss = criterion(logits, y)
        
        total_loss += loss.item() * X.size(0)
        total_correct += (logits.argmax(dim=1) == y).sum().item()
        total += X.size(0)
    
    return total_loss / total, total_correct / total

# Обучение
EPOCHS = 20
for epoch in range(EPOCHS):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    
    print(f"Epoch {epoch+1}/{EPOCHS}  "
          f"Train: loss={train_loss:.4f}, acc={train_acc:.4f}  "
          f"Val: loss={val_loss:.4f}, acc={val_acc:.4f}")

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

# Только weights (RECOMMENDED)
torch.save(model.state_dict(), "model.pt")

# Загрузка
model = MyModel(784, 256, 10)
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()

# Полный checkpoint (для resuming training)
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": loss,
}, "checkpoint.pt")

checkpoint = torch.load("checkpoint.pt")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

TorchScript — production export

# Tracing (нужны input examples)
example_input = torch.randn(1, 784).to(device)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")

# Scripting (full Python control flow)
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

# Загрузка (Python не нужен!)
loaded = torch.jit.load("model_traced.pt")
output = loaded(torch.randn(1, 784))

ONNX export

torch.onnx.export(
    model,
    example_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=17,
)

# Загрузка в ONNX Runtime
import onnxruntime as ort
sess = ort.InferenceSession("model.onnx")
output = sess.run(None, {"input": input_array})[0]

Интеграция с backend

PyTorch модель в FastAPI

from fastapi import FastAPI
from pydantic import BaseModel
import torch
from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app):
    app.state.device = "cuda" if torch.cuda.is_available() else "cpu"
    app.state.model = MyModel(784, 256, 10).to(app.state.device)
    app.state.model.load_state_dict(torch.load("model.pt", map_location=app.state.device))
    app.state.model.eval()
    yield

app = FastAPI(lifespan=lifespan)

class Input(BaseModel):
    features: list[float]  # 784 элемента

@app.post("/predict")
@torch.no_grad()
def predict(data: Input):
    X = torch.tensor([data.features], dtype=torch.float32).to(app.state.device)
    logits = app.state.model(X)
    probs = torch.softmax(logits, dim=1)
    pred_class = probs.argmax(dim=1).item()
    confidence = probs[0, pred_class].item()
    return {"class": pred_class, "confidence": confidence}

Batch prediction (эффективно)

@app.post("/predict/batch")
@torch.no_grad()
def predict_batch(items: list[Input]):
    X = torch.tensor([item.features for item in items], dtype=torch.float32).to(device)
    logits = app.state.model(X)
    probs = torch.softmax(logits, dim=1)
    return [
        {"class": int(p.argmax().item()), "confidence": float(p.max().item())}
        for p in probs
    ]

Production советы

1. model.eval() — Dropout и BatchNorm работают иначе в production
2. torch.no_grad() — gradient tracking отключается (быстрее + меньше memory)
3. torch.inference_mode() — no_grad + дополнительная оптимизация
4. Batching — 64 input в одном запросе — GPU используется лучше
5. TorchServe — в production: batching, versioning, A/B test (Месяц 6)
6. Async serving — asyncio + to_thread (CPU bound) или Triton/BentoML

Ресурсы

• PyTorch tutorials — pytorch.org/tutorials
• “Deep Learning with PyTorch” — Eli Stevens (free PDF: pytorch.org/deep-learning-with-pytorch)
• PyTorch Lightning — wrapper для уменьшения boilerplate
• Karpathy — “Let’s build GPT”(YouTube) — глубокий PyTorch
• Hugging Face Course — для PyTorch transformer’ов

🏋️ Упражнения

🟢 Easy

1. Создайте torch.randn(3, 4) tensor, выполните transpose, sum, mean.
2. С requires_grad=True найдите gradient f(x) = x³ при x=3.
3. Создайте nn.Linear(10, 1), выполните forward pass, выведите число параметров.

🟡 Medium

1. MNIST MLP: получите 95%+ accuracy на MNIST через 2-layer MLP.
2. Custom dataset: создайте свой Dataset class с CSV.
3. GPU check: тренируйте модель на CPU и GPU, измерьте разницу во времени.

🔴 Hard

1. FastAPI + PyTorch service: MNIST classifier, image upload, возвращает prediction. С Docker.
2. TorchScript benchmark: сравните latency обычной модели и TorchScript-версии (timeit).
3. Multi-GPU: train на 2+ GPU через nn.DataParallel или DistributedDataParallel (с Colab Pro или Kaggle).

Capstone

notebooks/month-03/02_pytorch_mnist.ipynb:

• Загрузите MNIST dataset через torchvision.datasets
• Напишите 3-layer MLP
• Train + Validation loop
• 97%+ accuracy на test set
• Confusion matrix
• Визуализируйте худшие примеры
• Экспортируйте модель в TorchScript
• Создайте FastAPI endpoint

✅ Чек-лист

• Создание tensor, операции, перенос на device
• Основы autograd (requires_grad, backward, grad)
• Subclassing nn.Module
• Batch loading через DataLoader
• Написание training loop (train + eval mode, zero_grad, optimizer.step)
• Сохранение и загрузка модели (state_dict)
• Export TorchScript или ONNX
• PyTorch serving в FastAPI

Переходим к TensorFlow и Keras.