Optimization & Regularization

Обратная связь

Хотелось бы больше погружения в код, как что работает в коде

Останавливайте, спрашивайте 🙏🏻

---

Обратная связь по обратной связи

• Мало

---

Recap

🧮 Мотивация

Нейронная сеть – большая формула с коэффициентами (весами).

Нейросети сами выучивают признаки, которые важны для предикта. Не надо заниматься фича-инженирингом.

📉 Градиенты

Для вычисления градиентов весов мы используем алгоритм обратного распространения ошибки.

📏 Нелинейность

Если из нейросети убрать все нелинейные преобразования, она выродится в линейную регрессию. ReLU - неплохой бейзлайн для ф-и активации.

🔮 Инициализация

Веса нейросети по умолчанию инициализируются рандомно, но так, чтобы сохранить дисперсию на каждом слое.

---

Оптимизация нейросетей. Регуляризация.

План занятия

Цели

• Понимать, что делает оптимизатор (SGD / momentum / Adam) и какие у него затраты по памяти.
• Понимать базовые методы регуляризации: dropout, weight decay (Adam vs AdamW), label smoothing.
• Уметь пользоваться LR scheduler и warmup
• Логировать метрики в TensorBoard.

---

Функции потерь

Делаем регрессию

Mean Squared Error (дока PyTorch)

MSE (Mean Squared Error) — регрессия: \[\mathcal{L}_{MSE} = \frac{1}{n}\sum_{i=1}^{n}(y_i - \hat{y}_i)^2\]

import torch.nn as nn
import torch

loss = nn.MSELoss()
input = torch.randn(3, 2, requires_grad=True)
target = torch.randn(3, 2)

output = loss(input, target)
output

Делаем классификацию:

Cross Entropy (дока PyTorch)

CE (Cross Entropy) — многоклассовая классификация (\(C\) классов, \(\hat{y}_i\) — вероятность класса \(i\) после softmax): \[\mathcal{L}_{CE} = -\sum_{i=1}^{C} y_i \log(\hat{y}_i)\]

# Example of target with class indices

import torch.nn as nn

loss = nn.CrossEntropyLoss() # nn.Softmax + nn.NLLoss
input = torch.randn(3, 5, requires_grad=True) # [ batch_size, class_probability ]

# [ batch_size ] -- метки каждого класса, каждая метка -- это число от [0 до 5)
target = torch.empty(3, dtype=torch.long).random_(5)

output = loss(input, target)
output

❓ Вопрос: зачем делать Softmax? Ведь на выходе модельки мы уже получаем логиты?

Делаем теггирование:

Binary Cross Entropy (дока PyTorch)

BCE (Binary Cross Entropy) — бинарная классификация. По элементам батча (\(N\) — размер батча, \(x_n\) — предсказание, \(y_n\) — целевой класс): \[\ell(x, y) = L = (l_1, \dots, l_N)^\top, \quad l_n = -w_n \left[ y_n \log x_n + (1 - y_n) \log(1 - x_n) \right]\] (веса \(w_n\) по умолчанию равны 1). Если reduction не 'none': \[\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{reduction} = \text{`mean';} \\ \operatorname{sum}(L), & \text{reduction} = \text{`sum'.} \end{cases}\]

Примером может быть задача тэгорования треков: один трек может именть несколько жанров: pop, rock, jazz, russian, 1980s…

sigmoid = nn.Sigmoid()
loss = nn.BCELoss()

input = torch.randn(3, requires_grad=True)
target = torch.empty(3).random_(2)

output = loss(sigmoid(input), target)
output

❓ Вопрос: Чем отличается от классификации?

Оптимизация

Моделька

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self, dropout_p=0.5):
        super().__init__()

        self.inner = nn.Sequential(nn.Linear(2, 16), # заведомо довольно большая модель, чтобы очевиднее был эффект переобучения и дропаута
                                   nn.ReLU(),
                                   nn.Linear(16, 16),
                                   nn.ReLU(),
                                   nn.Linear(16, 1))

    def forward(self, x):
        return self.inner(x)

# В этом методе градиенты всегда будут рандомными!
# В реальности данные и метки должны браться из датасета, а не из рандома
def example_compute_random_gradients(model):
    batch_size = 3
    model_input = torch.rand([batch_size, 2])
    target_batch_values = torch.randint(0, 1, [batch_size]) * 2 - 1 # -1/1 target values

    model_prediction = model.forward(model_input)

    loss = F.relu(1 - target_batch_values * model_prediction).mean()

    model.zero_grad()
    loss.backward()  # compute loss

Hands optimizer

model = MLP()
example_compute_random_gradients(model)

learning_rate = 0.01

# optimization step
for p in model.parameters():
    p.data = p.data - learning_rate * p.grad

Pytorch Optimizer

from torch.optim import SGD, Adam

model = MLP()
sgd_optimizer = SGD(model.parameters(), lr=learning_rate)

example_compute_random_gradients(model)

sgd_optimizer.step() # итериуется по сохраненным на этапе инициализации параметрам и обновляет параметры

## Другие оптимизаторы
from torch.optim import Adam

adam_optimizer = Adam(model.parameters(), lr=learning_rate)

example_compute_random_gradients(model)

adam_optimizer.step()

Adam: формулы и интуиция

Обозначим градиент на шаге \(t\): \(g_t = \nabla_\theta \mathcal{L}(\theta_t)\).

Экспоненциальные скользящие средние (моменты): - Первый момент (средний градиент):

\[ m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \]

• Второй момент (средний квадрат градиента):

\[ v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \]

Bias correction (важно на первых шагах, потому что \(m_0=v_0=0\)):

\[ \hat m_t = \frac{m_t}{1-\beta_1^t}, \quad \hat v_t = \frac{v_t}{1-\beta_2^t} \]

Шаг обновления:

\[ \theta_{t+1} = \theta_t - \alpha \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon} \]

Интуиция: - \(\hat m_t\) похож на momentum (направление шага) - \(\sqrt{\hat v_t}\) нормирует шаг по координатам (где градиенты «шумнее/больше» — там шаг меньше) - поэтому Adam часто быстрее «разгоняется» и устойчивее к плохому масштабу признаков/слоев

Можем вспомнить, что оптимизатор хранит статиситики по градиентам параметров (моменты)!

Для хранения моментов тоже используется память на видео-карте.

adam_optimizer.state_dict()

Кстати! А как работает скользящее среднее?

import matplotlib.pyplot as plt

timeseries_len = 100

outlier = torch.zeros(timeseries_len)
outlier[55] = 3

           # рандомный шум                   + квадратный корень                       + выброс
raw_data = torch.rand([timeseries_len]) + torch.sqrt(torch.arange(timeseries_len)) + outlier
raw_data = raw_data / 100

beta_values = torch.tensor([ 0.5, 0.9, 0.99 ])

mooving_avarages = torch.zeros([beta_values.shape[0], timeseries_len])
mooving_avarages[:, 0] = raw_data[0]

for i in range(1, timeseries_len):
    mooving_avarages[:, i] = mooving_avarages[:, i-1] * beta_values + (1-beta_values) * raw_data[i]

plt.plot(raw_data, label='raw')
for beta_i, beta in enumerate(beta_values.numpy().tolist()):
    plt.plot(mooving_avarages[beta_i, :], label=f"beta {beta:.2f}")

plt.legend()
plt.title("Mooving average demo")

❓ Вопрос: Какую проблему решает скользящее среднее?

Demo optimizers

!nvidia-smi

import torch
import torch.optim as optim

# Используем MLP из этого ноутбука
mlp_adam = MLP()
mlp_sgd = MLP()

mlp_adam

# Сколько памяти съедает состояние оптимизатора?
# (на GPU это тоже хранится, поэтому Adam «дороже» по памяти)

import torch


def optimizer_state_size_bytes(optimizer: torch.optim.Optimizer) -> int:
    total = 0
    for state in optimizer.state.values():
        if not isinstance(state, dict):
            continue
        for v in state.values():
            if torch.is_tensor(v):
                total += v.numel() * v.element_size()
    return total


def model_params_size_bytes(model: torch.nn.Module) -> int:
    return sum(p.numel() * p.element_size() for p in model.parameters() if p.requires_grad)


def init_optimizer_state(optimizer: torch.optim.Optimizer, model: torch.nn.Module) -> None:
    x = torch.randn(2048, 2)
    y = torch.randn(2048, 1)

    optimizer.zero_grad()
    pred = model(x)
    loss = torch.nn.functional.mse_loss(pred, y)
    loss.backward()
    optimizer.step()


model_sgd = MLP()
sgd = torch.optim.SGD(model_sgd.parameters(), lr=1e-2)
init_optimizer_state(sgd, model_sgd)

model_sgd_momentum = MLP()
sgd_momentum = torch.optim.SGD(model_sgd_momentum.parameters(), lr=1e-2, momentum=0.9)
init_optimizer_state(sgd_momentum, model_sgd_momentum)

model_adam = MLP()
adam = torch.optim.Adam(model_adam.parameters(), lr=1e-3)
init_optimizer_state(adam, model_adam)

print("MLP params:", round(model_params_size_bytes(model_sgd) / 1024, 2), "KB")
print("SGD state:", round(optimizer_state_size_bytes(sgd) / 1024, 2), "KB")
print("SGD+momentum state:", round(optimizer_state_size_bytes(sgd_momentum) / 1024, 2), "KB")
print("Adam state:", round(optimizer_state_size_bytes(adam) / 1024, 2), "KB")

Регуляризация

• Dropout
• Weight Decay
• Label Smoothing

import torch
import torch.nn as nn

Dropout

Dropout

При обучении каждый элемент входа с вероятностью \(p\) обнуляется, затем результат масштабируется на \(1/(1-p)\) (чтобы сохранить мат. ожидание): \[y_i = \frac{x_i \cdot m_i}{1 - p}, \quad m_i \sim \operatorname{Bernoulli}(1-p)\] В режиме eval: \(y = x\) (dropout отключён).

dropout = nn.Dropout(p=0.5)
dropout.eval()
print(dropout.training) # важный флажок для дропаута - тк дропаут имеет разное поведение в процессе обучения и в процессе вычисления сети
dropout.train()
print(dropout.training)

dropout

list(dropout.parameters())

t = torch.rand([3, 5], dtype=torch.float32)

print("t.mean()", t.mean().item())

t

dropout.eval()
t_dropouted = dropout(t)

print("t_dropouted.mean()", t_dropouted.mean().item())

assert torch.allclose(t_dropouted, t), "eval dropout module do nothing"

t_dropouted

dropout.train()
t_dropouted = dropout(t)

print("t_dropouted.mean()", t_dropouted.mean().item())

print("zeros count", (t_dropouted == 0.0).sum().item(), "of", t_dropouted.numel())

t_dropouted

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.dropout = nn.Dropout()

m = Model()
m.training

m = m.eval()
m.dropout.training

m = m.train()
m.dropout.training

LabelSmoothing

nn.CrossEntropyLoss(label_smoothing=0.1)

WeightDecay

• Adam (…, weight_decay=0.1)

• AdamW

Stack Overflow: AdamW and Adam with weight decay

В чем разница:

Adam + weight_decay - регуляризация идет через градиенты

final_loss = loss + wd * all_weights.pow(2).sum() / 2
# backward ...
# optimizer.step:
w = w - lr * w.grad # grad contains regularisation

AdamW - прямое обновление весов

w = w - lr * w.grad - lr * wd * w

Аспект	Adam + weight_decay	AdamW
Функция потерь	L_final = L + λ/2 Σw²	L_final = L
Обновление весов	w ← w - α·m/√(v+ε)	w ← w - α·m/√(v+ε) - α·λ·w
Сила регуляризации	Различается для разных параметров	Одинаковая для всех параметров
Зависимость от градиентов	λ попадает под адаптивный lr	λ не зависит от градиентов
Для SGD	✅ Разницы нет!	✅ Разницы нет!
Применение	Устаревший подход	Современный стандарт для трансформеров, лучше работает

import torch.optim as optim

dummy_module_params = nn.Linear(1, 1).parameters()
adam = optim.Adam(dummy_module_params, weight_decay=0.01)

dummy_module_params = nn.Linear(1, 1).parameters()
adamw = optim.AdamW(dummy_module_params, weight_decay=0.1)

Обновим Trainer и нашу модель

• Добавим поддержку lr_sheduler
• В модельку добавим дропаут

from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from torchvision.datasets import FashionMNIST
import matplotlib.pyplot as plt
import torch.nn as nn
import numpy as np

from torchvision.transforms import ToTensor

from torch.optim import Adam

# ничего нового, копипаста с предыдущего сема
class FMNISTImageSet:
    def __init__(self, train=True, transform=None):
        self.data = FashionMNIST("./tmp", train=train, download=True)

        self.transform = transform

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

    def __getitem__(self, item):
        # сделайте одноканальную картинку [1, 28, 28] с float32
        sample, label = self.data[item]

        if self.transform is not None:
            sample = self.transform(sample)
        else:
            sample = np.array(sample, dtype=np.float32)[None:, ...] / 255


        return dict(
            sample=sample,
            label=label,
        )

fmnist_train = FMNISTImageSet(train=True, transform=ToTensor())
fmnist_val = FMNISTImageSet(train=False, transform=ToTensor())

# Новый модуль
flatten_module = nn.Flatten(start_dim=1, end_dim=-1)
flatten_module(torch.rand([ 64, 28, 28 ])).shape

# Что добавилось:
# * `nn.Dropout`, и параметр `dropout_p` - новый слой и параметр в `__init__`
# * `nn.Flatten` - новый слой для схлопывания размерностей на входящих тензорах
# * `label_smoothing=0.0` - новый параметр

class MLPModel(nn.Module):
    """
    Реализация MLP. Похожее мы уже видели на 1 и 2 семинаре.
    """

    def __init__(self, dropout_p=0.5, activation=None, label_smoothing=0.0):
        """
        Инициализация модели.

        :param activation: Функция активации (по умолчанию ReLU).
        :param dropout_p: Процент зануляемых значений в активациях
        """
        super().__init__()

        self.label_smoothing = label_smoothing

        if activation is None:
            activation = nn.ReLU()

        self.inner = nn.Sequential(
            # input ~ [ bs, 28, 28 ]
            nn.Flatten(), # ~ [ bs, 28 * 28 ]
            nn.Linear(784, 100),
            nn.Dropout(p=dropout_p),
            activation,
            nn.Linear(100, 10),
        )

        # nn.Sequential последовательно применяет переданные слои
        # self.linear_1 = nn.Linear(784, 100)
        # self.activation = nn.ReLU()
        # self.linear_2 = nn.Linear(100, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Прямой проход данных через модель.

        :param x: Входные данные ~ [ bs, 784 ]
        :return: Выходные логиты модели ~ [ bs, 10 ]
        """
        return self.inner(x)

    def compute_all(self, batch):
        """
        Вычисляет лосс и точность по батчу данных.

        :param batch: Батч данных, содержащий 'sample' (входные данные) и 'label' (метки классов).
        :return: Значение функции потерь и словарь метрик.
        """

        x = batch['sample']    # [ bs, 784 ]
        y = batch['label']     # [ bs ] torch.LongTensor
        logits = self.inner(x) # [ bs, 10 ]

        loss = F.cross_entropy(logits, y, label_smoothing=self.label_smoothing)
        acc = (logits.argmax(axis=1) == y).float().mean().cpu().numpy()
        metrics = {
            "metrics/acc": acc,
            "metrics/loss": loss.detach().item(),
        }
        return loss, metrics

Про Trainer

Мы уже писали похожий Trainer на прошлом семинаре. Здесь он чуть расширен: - Поддержка TrainerConfig - класс, в котором описываем конфигурацию обучения - Хуки: выбираем, когда делаем lr_scheduler.step() (per_batch или per_epoch)

Важно: сам Trainer — это просто удобная обертка над циклом обучения, чтобы не копипастить одно и то же.

# TensorBoard (Colab)
# Логи будут писаться в папку runs/03_seminar

%load_ext tensorboard

LOG_DIR = "runs/03_seminar"

# Запуск интерфейса можно делать после первого обучения:
# %tensorboard --logdir runs/03_seminar

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler, ReduceLROnPlateau
from torch.utils.tensorboard import SummaryWriter


@dataclass(frozen=True)
class TrainerConfig:
    log_dir: str = "runs/03_seminar"
    batch_size: int = 128
    lr_scheduler_type: str | None = None  # None | 'per_batch' | 'per_epoch'


class Trainer:
    """Обертка над обучением модели.

    Держит вместе:
    - model
    - optimizer
    - train/val datasets + dataloaders
    - (опционально) lr scheduler
    - логирование в TensorBoard
    """

    def __init__(
        self,
        model: nn.Module,
        optimizer: Optimizer,
        train_dataset: Iterable,
        val_dataset: Iterable,
        *,
        config: TrainerConfig | None = None,
        lr_scheduler: LRScheduler | None = None,
    ):
        if config is None:
            config = TrainerConfig()

        if config.lr_scheduler_type not in [None, "per_batch", "per_epoch"]:
            raise ValueError(
                "lr_scheduler_type must be one of: None, 'per_batch', 'per_epoch'. "
                f"Not: {config.lr_scheduler_type}"
            )

        self.model = model
        self.optimizer = optimizer
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset

        self.batch_size = config.batch_size
        self.lr_scheduler = lr_scheduler
        self.lr_scheduler_type = config.lr_scheduler_type

        self.device: str | int = "cpu"
        if torch.cuda.is_available():
            self.device = torch.cuda.current_device()
            self.model = self.model.to(self.device)

        self.global_step = 0
        self.writer = SummaryWriter(log_dir=config.log_dir)

    def close(self) -> None:
        self.writer.flush()
        self.writer.close()

    def log_scalars(self, log_data: dict[str, float], step: int | None = None) -> None:
        if step is None:
            step = self.global_step

        for key, value in log_data.items():
            self.writer.add_scalar(key, float(value), step)

    def save_checkpoint(self, path: str) -> None:
        """Сохраняет веса модели в указанный путь."""
        torch.save(self.model.state_dict(), path)

    def train(self, num_epochs: int) -> None:
        train_loader = DataLoader(
            self.train_dataset,
            shuffle=True,
            pin_memory=True,
            batch_size=self.batch_size,
        )

        # Валидационный датасет не стоит перемешивать
        val_loader = DataLoader(
            self.val_dataset,
            shuffle=False,
            pin_memory=True,
            batch_size=self.batch_size,
        )

        best_loss = float("inf")

        for epoch in tqdm(range(num_epochs)):
            self.model.train()
            for batch in train_loader:
                self.training_step(batch)
                self.post_train_batch()
                self.global_step += 1

            self.model.eval()

            val_losses: list[float] = []
            for batch in val_loader:
                loss = self.validation_step(batch)
                val_losses.append(float(loss.item()))

            val_loss = float(np.mean(val_losses))
            self.log_scalars({"validation/loss": val_loss}, step=epoch)

            # Сохраняем только наилучшую модель
            if val_loss < best_loss:
                self.save_checkpoint("./best_checkpoint.pth")
                best_loss = val_loss

            self.post_val_stage(val_loss)

        self.close()

    def training_step(self, batch) -> None:
        self.optimizer.zero_grad()

        batch = {k: v.to(self.device) for k, v in batch.items()}
        loss, details = self.model.compute_all(batch)

        loss.backward()
        self.optimizer.step()

        # Логирование метрик обучения
        log_metrics = {f"train/{k}": float(v) for k, v in details.items()}
        log_metrics["train/loss"] = float(loss.item())
        self.log_scalars(log_metrics)

    def validation_step(self, batch):
        batch = {k: v.to(self.device) for k, v in batch.items()}
        loss, _details = self.model.compute_all(batch)
        return loss

    def post_train_batch(self) -> None:
        # called after every train batch
        if self.lr_scheduler is None or self.lr_scheduler_type != "per_batch":
            return

        self.lr_scheduler.step()

        # Берем lr для первой группы параметров
        lr = float(self.optimizer.param_groups[0]["lr"])
        self.log_scalars({"train/learning_rate": lr})

    def post_val_stage(self, val_loss: float) -> None:
        if self.lr_scheduler is None or self.lr_scheduler_type != "per_epoch":
            return

        if isinstance(self.lr_scheduler, ReduceLROnPlateau):
            self.lr_scheduler.step(val_loss)
        else:
            self.lr_scheduler.step()

        lr = float(self.optimizer.param_groups[0]["lr"])
        self.log_scalars({"validation/learning_rate": lr}, step=self.global_step)

Не забудь поменять Runtime Type в колабе на GPU!

import torch.utils.data

# берем часть исходной обучающей выборки для большей наглядности эффекта применения дропаута
# на большом датасете сложнее переобучиться, поэтому мы искусственно уменьшаем его для наглядности на занятии
# на практике не надо так делать!
fmnist_train_subset = torch.utils.data.Subset(fmnist_train, range(0, len(fmnist_train), 400))

dropout=0

model = MLPModel(dropout_p=0.0)
optimizer = Adam(model.parameters(), lr=3e-3)

trainer = Trainer(
    model,
    optimizer,
    fmnist_train_subset,
    fmnist_val,
    config=TrainerConfig(log_dir=f"{LOG_DIR}/dropout_p_0.0"),
)
trainer.train(30)

dropout=0.3

model = MLPModel(dropout_p=0.3)
optimizer = Adam(model.parameters(), lr=3e-3)

trainer = Trainer(
    model,
    optimizer,
    fmnist_train_subset,
    fmnist_val,
    config=TrainerConfig(log_dir=f"{LOG_DIR}/dropout_p_0.3"),
)
trainer.train(30)

dropout=0.3, label_smoothing=0.1

model = MLPModel(dropout_p=0.3, label_smoothing=0.1)
optimizer = Adam(model.parameters(), lr=3e-3)

trainer = Trainer(
    model,
    optimizer,
    fmnist_train_subset,
    fmnist_val,
    config=TrainerConfig(log_dir=f"{LOG_DIR}/dropout_p_0.3_label_smoothing_0.1"),
)
trainer.train(30)

TensorBoard (как смотреть логи в Colab)

1. Убедитесь, что вы запускали обучение хотя бы один раз (в runs/03_seminar появятся event-файлы).
2. Запустите TensorBoard в отдельной ячейке:

%tensorboard --logdir runs/03_seminar

Если вы запускаете несколько экспериментов, удобно логировать в подпапки: - runs/03_seminar/dropout_p_0.0 - runs/03_seminar/dropout_p_0.3 - runs/03_seminar/CyclicLR

TensorBoard сам подхватит все подпапки.

%tensorboard --logdir runs/03_seminar

LR Scheduler

Два типа расписаний:

• по эпохам (StepLR, ReduceLROnPlateau, …)

scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
for epoch in range(epochs):
    train(...)
    validate(...)
    scheduler.step()

• по батчам (Cyclic, 1cycle, …)

scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.01, max_lr=0.1)
for epoch in range(epochs):
    # train(...)
    for batch in data_loader:
        train_batch(...)
        scheduler.step()
    # validate(...)

WarmUp

Проблема: очень большие сетки может сильно разнести, если сразу начать их оптимизировать на больших батчах.

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

import math
import matplotlib.pyplot as plt
import torch


def plot_lr_schedule(optimizer, scheduler, num_steps: int, title: str, step_size=1) -> None:
    lr_history = []
    steps_history = []
    for i in range(0, num_steps, step_size):
        steps_history.append(i)
        lr_history.append(float(optimizer.param_groups[0]["lr"]))
        scheduler.step()

    plt.plot(steps_history, lr_history)
    plt.title(title)
    plt.xlabel("step")
    plt.ylabel("lr")
    plt.grid(True)


# Визуализация warmup + cosine (через LambdaLR)
base_lr = 3e-4
num_steps = 200
warmup_steps = 20

param = torch.nn.Parameter(torch.zeros(()))
optimizer = torch.optim.Adam([param], lr=base_lr)

# lr_factor(step): [0..1] на warmup, затем cosine до 0

def lr_factor(step: int) -> float:
    if step < warmup_steps:
        return float(step + 1) / float(warmup_steps)

    progress = float(step - warmup_steps) / float(max(1, num_steps - warmup_steps))
    return 0.5 * (1.0 + math.cos(math.pi * progress))


scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_factor)

step_optimizer = torch.optim.Adam([param], lr=base_lr)
step_lr_scheduler = torch.optim.lr_scheduler.StepLR(step_optimizer, step_size=30, gamma=0.999)

plt.figure(figsize=(8, 3))
plot_lr_schedule(optimizer, scheduler, num_steps=num_steps, title="Warmup + Cosine (LambdaLR)")
plt.show()

Обучим модель с шедулером и warmup’ом

model = MLPModel(dropout_p=0.3, label_smoothing=0.1)

base_lr = 3e-3
optimizer = Adam(model.parameters(), lr=base_lr)

# TODO avoid hardcode, build scheduler in trainer
num_steps = 59
warmup_steps = 20

optimizer = torch.optim.Adam(model.parameters(), lr=base_lr)

# lr_factor(step): [0..1] на warmup, затем cosine до 0
def lr_factor(step: int) -> float:
    if step < warmup_steps:
        return float(step + 1) / float(warmup_steps)

    progress = float(step - warmup_steps) / float(max(1, num_steps - warmup_steps))
    return 0.5 * (1.0 + math.cos(math.pi * progress))


scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_factor)


trainer = Trainer(
    model,
    optimizer,
    fmnist_train_subset,
    fmnist_val,
    config=TrainerConfig(
        log_dir=f"{LOG_DIR}/cosine_with_warmup",
        lr_scheduler_type="per_batch",
    ),
    lr_scheduler=scheduler,
)
trainer.train(30)

Оптимизаторы

Бонусная домашка по какой-то из тем (голосуем): - Schedule-free - Muon

Тюнинг LR

Проблема: с масштабированием моделей нужно с нуля подбирать lr и другие гиперпараметры.

Решение: \(\mu P\) параметризация или (вроде) Muon

⚠️ Pytorch реализация SGD with momentum

Обновление параметра: \(\theta_t = \theta_{t-1} - \eta \cdot v_t\) (в обоих вариантах). Различие — в правиле накопления буфера \(v_t\).

Возможно вы ожидали, что

\[v_t = (1 - \mu) \, g_t + \mu \, v_{t-1}\] (\(g_t\) — градиент на шаге \(t\), \(\mu\) — momentum, экспоненциальное сглаживание градиентов)

self.param_momentum[p] = (1 - self.momentum) * p.grad + self.momentum * self.param_momentum[p]

Но pytorch implementation

\[v_t = g_t + \mu \, v_{t-1}\] (накопленная сумма градиентов с затуханием \(\mu\))

self.param_momentum[p] = p.grad + self.momentum * self.param_momentum[p]

❓ Что означает на практике?

Когда хотите использовать SGD w/ momentum, умножайте learning rate на 1/(1-momentum), если хотите переиспользовать LR из SGD.

StackOverflow

---

Резюме

Блиц

Может ли быть лосс отрицательным?

Зачем нужен lr warmup?

Зачем нужен lr scheduling?

Как нужно изменить lr, если мы увеличили batch_size в 3 раза?

Как выбрать хорошие начальные значения для lr?

Как будет обучаться модель, если в нее добавить слой nn.Dropout(p=1.0)

Зачем внутри Dropout делается умножение на 1/(1-p)? Что будет, если этого не сделать?

Как изменяется поведение модуля nn.Dropout после того, как мы переключили модель в режим вычисления (model.eval())?

Для чего нужен nn.Dropout?

Мы хотим обучить модель на 1M параметров с помощью SGD и с помощью Adam. Как изменится необходимый объем видеопамяти для обучения модели в зависимости от оптимизатора?

Зачем Adam использует скользящее среднее для вычисления моментов?

Чем Adam + weight decay отличается от AdamW?

Нашел отзыв на яндекс маркете. Про какой это оптимизатор?

Домашки:

• dropout
• optimization
• [бонусная?] [WIP] muon
• [бонусная?] [WIP] scheduler free adam