{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Обучение нейросетей: Продвинутые приемы" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "import albumentations as A\n", "import albumentations.pytorch.transforms\n", "import lightning as L\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import nvidia.dali as dali\n", "import nvidia.dali.fn as fn\n", "import nvidia.dali.plugin.pytorch as dali_pytorch\n", "import nvidia.dali.types as types\n", "import timm\n", "import torch\n", "import torch.nn.functional as F\n", "import torchmetrics\n", "import torchvision\n", "from nvidia.dali import pipeline_def\n", "from torch import nn\n", "from torch.utils import data\n", "from tqdm.auto import tqdm\n", "\n", "# Increase these if figures appear small\n", "plt.rcParams[\"figure.figsize\"] = fx, fy = (14.08, 6.40)\n", "\n", "# Not using `bfloat16` matrix multiplication for consistency\n", "# You might get better performance without much precision loss\n", "# by setting this to \"medium\" on some devices\n", "torch.set_float32_matmul_precision(\"high\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def show_images(images, titles=[]):\n", " num = len(images)\n", " fig, axs = plt.subplots(nrows=1, ncols=num, squeeze=True, layout=\"constrained\")\n", " axs = [axs] if num <= 1 else axs\n", " for i in range(num):\n", " ax = axs[i]\n", " ax.imshow(images[i])\n", " ax.axis(\"off\")\n", " if titles != []:\n", " ax.set_title(titles[i])\n", "\n", " plt.show(fig)\n", " plt.close(fig)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "MEAN = (0.4914, 0.4822, 0.4465)\n", "STD = (0.2023, 0.1994, 0.2010)\n", "\n", "BATCH_SIZE = 32\n", "MAX_EPOCHS = 10\n", "BASE_LR = 0.05\n", "NUM_CLASSES = 10\n", "\n", "NUM_WORKERS = os.cpu_count()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Данные" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "В качестве датасета будем использовать CIFAR10" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def CIFAR10(**kwargs):\n", " kwargs.setdefault(\"root\", \"cifar10\")\n", " return torchvision.datasets.CIFAR10(**kwargs)\n", "\n", "\n", "# Download the dataset\n", "CIFAR10(download=True);" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "dataset = CIFAR10(transform=None)\n", "\n", "images = []\n", "titles = []\n", "rows, cols = 4, 8\n", "for idx in range(rows * cols):\n", " image, label = dataset[idx]\n", " images.append(image)\n", " titles.append(f\"Class {label} ({dataset.classes[label]})\")\n", "\n", "for _ in range(rows):\n", " show_images(\n", " [images.pop() for _ in range(cols)],\n", " [titles.pop() for _ in range(cols)],\n", " )" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "augmentations = [\n", " A.Affine(translate_percent=0.05, scale=(1.0, 1.05), rotate=15, p=0.5),\n", " A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.5),\n", " A.RandomBrightnessContrast(p=0.5),\n", "]\n", "common_transforms = [\n", " A.Normalize(mean=MEAN, std=STD),\n", " A.pytorch.transforms.ToTensorV2(),\n", "]\n", "\n", "MyTrainTransform = A.Compose(augmentations + common_transforms)\n", "MyValidTransform = A.Compose(common_transforms)\n", "\n", "\n", "def my_train_transform(image):\n", " return MyTrainTransform(image=np.array(image))[\"image\"]\n", "\n", "\n", "def my_valid_transform(image):\n", " return MyValidTransform(image=np.array(image))[\"image\"]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# NOTE: technically, CIFAR10(train=False) is the TEST set,\n", "# not the validation set, but we use it here for simplicity\n", "\n", "\n", "ds_train = CIFAR10(transform=my_train_transform, train=True)\n", "ds_valid = CIFAR10(transform=my_valid_transform, train=False)\n", "\n", "dl_train = data.DataLoader(\n", " ds_train,\n", " batch_size=BATCH_SIZE,\n", " shuffle=True,\n", " drop_last=True,\n", " num_workers=NUM_WORKERS,\n", ")\n", "dl_valid = data.DataLoader(\n", " ds_valid,\n", " batch_size=BATCH_SIZE,\n", " shuffle=False,\n", " drop_last=False,\n", " num_workers=NUM_WORKERS,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Загрузка готовых моделей\n", "\n", "В качестве базовой модели возьмем ResNet-18. Есть несколько способов загрузить готовую архитектуру и веса." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Torchvision Models\n", "\n", "Список моделей можно посмотреть в [документации](https://pytorch.org/vision/stable/models.html)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_resnet_torchvision(num_classes, transfer=True):\n", " weights = torchvision.models.ResNet18_Weights.DEFAULT if transfer else None\n", " model = torchvision.models.resnet18(weights=weights)\n", "\n", " linear_size = model.fc.in_features\n", " model.fc = nn.Linear(linear_size, num_classes)\n", "\n", " return model" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Pytorch Image Models (aka `timm`)\n", "\n", "Официальный репозиторий: [pytorch-image-models](https://github.com/huggingface/pytorch-image-models).\n", "\n", "Выведем список моделей `resnet` из `timm`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "timm.list_models(filter=\"*resnet*18*\", pretrained=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "В списке одна и та же модель может появляться несколько раз с разными суффиксами. Эти суффиксы соответствуют разным процедурам обучения — например, `.a1_in1k` означает, что модель обучалась на датасете ImageNet-1k по схеме A1 из статьи [ResNet strikes back](https://arxiv.org/abs/2110.00476).\n", "\n", "Замеры метрик, скоростей работы, потребления памяти и т.п. для различных моделей на различных девайсах можно посмотреть [в папке `results` в репозитории `timm`](https://github.com/huggingface/pytorch-image-models/tree/main/results)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_resnet_timm(num_classes, transfer=True):\n", " return timm.create_model(\n", " \"resnet18.a1_in1k\",\n", " pretrained=transfer,\n", " num_classes=num_classes,\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## PyTorch Lightning" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Почему PyTorch Lightning?\n", "\n", "#### Много реализованного функционала\n", "\n", "* Обучение на GPU/нескольких GPU/TPU \n", "\n", "* Multi-node training\n", "\n", "* Авто логгирование\n", "\n", "* Gradient accumulation\n", "\n", "* …\n", "\n", "#### Меньше boilerplate кода\n", "\n", "Любое исследование можно начать с простого кода и быстро добавить более сложные (уже реализованные) вещи, такие как: обучение на GPU, 16-bit, сохрание чекпоинтов, логгирование, …\n", "\n", "PyTorch Lightning сделал и протестировал много всего за Вас, чтобы Вы могли сосредоточиться на исследовании других идей.\n", "\n", "Меньше инженерного кода означает:\n", "\n", "* меньше багов\n", "\n", "* более быстрое итерирование и прототипирование\n", "\n", "#### Стандартная структура кода\n", "\n", "PyTorch Lightning задает дефолтную \"opinionated\" схему организации кода. Переходя в новый проект не нужно с нуля разбираться в том, как авторы переизобрели велосипед. Callback функции имеют предсказуемые имена, меньше глобальных объектов, меньше магических констант и тп." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import lightning as L\n", "\n", "\n", "# Almost direct translation of the train function from the previous seminar\n", "class MyCIFAR10TrainingModule(L.LightningModule):\n", " def __init__(self):\n", " super().__init__()\n", " self.model = get_resnet_torchvision(num_classes=NUM_CLASSES)\n", " self.loss_fn = torch.nn.CrossEntropyLoss()\n", "\n", " self.train_loss = []\n", " self.valid_accs = []\n", "\n", " def configure_optimizers(self):\n", " return torch.optim.Adam(self.model.parameters(), lr=BASE_LR)\n", "\n", " def training_step(self, batch):\n", " x, y = batch\n", " p = self.model(x)\n", " loss = self.loss_fn(p, y)\n", " self.train_loss.append(loss.detach())\n", " return loss\n", "\n", " def validation_step(self, batch):\n", " x, y = batch\n", " p = self.model(x)\n", " matches = p.argmax(axis=-1) == y\n", " self.valid_accs.append(matches)\n", " return {}\n", "\n", " # The on_*_epoch_end callbacks below show that you CAN translate\n", " # your normal train loop almost directly into Lightning.\n", "\n", " # However, this is NOT the recommended / idiomatic way\n", " # to compute metrics, when using Lightning.\n", "\n", " def on_train_epoch_end(self):\n", " epoch_loss = torch.stack(self.train_loss).mean()\n", " print(\n", " f\"Epoch {self.trainer.current_epoch},\",\n", " f\"train_loss: {epoch_loss.item():.8f}\",\n", " )\n", " # don't forget to clear the saved losses\n", " self.train_loss.clear()\n", "\n", " def on_validation_epoch_end(self):\n", " epoch_accs = torch.concat(self.valid_accs).float().mean()\n", " print(\n", " f\"Epoch {self.trainer.current_epoch},\",\n", " f\"valid_accs: {epoch_accs.item():.8f}\",\n", " )\n", " # don't forget to clear the saved accuracies\n", " self.valid_accs.clear()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Чуть более идиоматичный LightningModule" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torchmetrics\n", "\n", "\n", "class MyCIFAR10TrainingModule(L.LightningModule):\n", " def __init__(\n", " self,\n", " *,\n", " # You can add more hyperparameters here\n", " num_classes=NUM_CLASSES,\n", " lr=BASE_LR,\n", " **kwargs,\n", " ):\n", " super().__init__(**kwargs)\n", " self.num_classes = num_classes\n", " self.lr = lr\n", "\n", " self.model = self.get_model()\n", " self.loss_fn = nn.CrossEntropyLoss()\n", "\n", " # Lightning can automatically handle metrics for you,\n", " # including tricky implementation details like accumulating\n", " # non-trivial metrics with multi-device reductions etc.\n", "\n", " # `torchmetrics` has a bunch of metrics already implemented, but\n", " # you can add your own by inheriting from `torchmetrics.Metric`\n", " # and implementing the required methods.\n", " self.accuracy = torchmetrics.classification.Accuracy(\n", " task=\"multiclass\",\n", " num_classes=self.num_classes,\n", " )\n", "\n", " def get_model(self):\n", " return get_resnet_torchvision(num_classes=10)\n", "\n", " def configure_optimizers(self):\n", " return torch.optim.Adam(self.model.parameters(), lr=self.lr)\n", "\n", " def training_step(self, batch):\n", " return self._step(batch, \"train\")\n", "\n", " def validation_step(self, batch):\n", " return self._step(batch, \"valid\")\n", "\n", " def _step(self, batch, kind):\n", " x, y = batch\n", " p = self.model(x)\n", "\n", " loss = self.loss_fn(p, y)\n", " accs = self.accuracy(p.argmax(axis=-1), y)\n", "\n", " return self._log_metrics(loss, accs, kind)\n", "\n", " # No manual on_*_epoch_end callbacks needed.\n", " # Lightning will automatically track the logged\n", " # metric objects and apply reductions as needed.\n", "\n", " def _log_metrics(self, loss, accs, kind):\n", " metrics = {}\n", " if loss is not None:\n", " metrics[f\"{kind}_loss\"] = loss\n", " if accs is not None:\n", " metrics[f\"{kind}_accs\"] = accs\n", " self.log_dict(\n", " metrics,\n", " prog_bar=True,\n", " logger=True,\n", " on_step=kind == \"train\",\n", " on_epoch=True,\n", " )\n", " return loss" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Init our training module\n", "training_module = MyCIFAR10TrainingModule()\n", "\n", "# Initialize a trainer\n", "trainer = L.Trainer(accelerator=\"auto\", default_root_dir=\"runs/basic\", max_epochs=MAX_EPOCHS)\n", "\n", "# Train the model ⚡\n", "trainer.fit(training_module, dl_train, dl_valid)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Логи и чекпоинты\n", "\n", "Как уже было сказано выше в Pytorch Lightning реализовано автоматическое логгирование.\n", "\n", "Сохранены все запуски: " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!tree runs/" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Как видите, Lightning автоматически сохранил чекпоинт последней эпохи. Если прервать обучение, то все равно будет сохранены результаты с последней успешно завершившейся эпохи.\n", "\n", "Обратите внимание, что чекпоинт - это не просто веса обученной модели. Например, в чекпоинт также включаются статистики оптимизатора." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "last_ckpt = \"runs/basic/lightning_logs/version_0/checkpoints/epoch=2-step=4686.ckpt\"\n", "ckpt = torch.load(last_ckpt, weights_only=True, map_location=\"cpu\")\n", "print(list(ckpt.keys()))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Таким образом, используя чекпоинт можно возобновить прерванный процесс обучения." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Continue training the model for 2 more epochs\n", "training_module = MyCIFAR10TrainingModule()\n", "trainer = L.Trainer(accelerator=\"auto\", default_root_dir=\"runs/basic_resume\", max_epochs=MAX_EPOCHS+2)\n", "trainer.fit(training_module, dl_train, dl_valid, ckpt_path=last_ckpt)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "training_module = MyCIFAR10TrainingModule.load_from_checkpoint(last_ckpt)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "training_module.model.layer1" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "sd = training_module.model.state_dict()\n", "torch.save(sd, \"model.pt\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Tensorboard\n", "\n", "Tensorboard - популярный инструментарий для визуализации различных объектов получаемых в процессе обучения. В Lightning по умолчанию включено автоматическое сохранения логов в формате Tensorboard. Кроме того, в Lightning есть поддержка для ряда других библитоек логирования (WandB, Comet, MLFlow, Neptune, простой CSV или даже свой кастомный логгер).\n", "\n", "Tensorboard можно запустить, выполнив комманду `tensorboard` в терминале и открыв распечатанную ссылку в браузере (обычно, http://localhost:6006/).\n", "\n", "Либо, используя специальное расширение для Jupyter, можно открыть Tensorboard как embedded окошечко для прямо в ноутбуке:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%load_ext tensorboard" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%tensorboard --logdir runs" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Кастомизация чекпоинтов\n", "\n", "Мы уже видели, что Pytorch Lightning сохраняет веса модели самостоятельно. Однако сохраняются веса только последней эпохи, метрики на которой могут быть не максимальны, чтобы это исправить самостоятельно напишем callback **`ModelCheckpoint`**.\n", "\n", "Также есть callback **`EarlyStopping`**, который позволяет автоматически останавливать обучение при отсутствие роста метрик в течение нескольких эпох." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def my_training_checkpoint():\n", " return L.pytorch.callbacks.ModelCheckpoint(\n", " filename=\"{epoch}-{valid_accs:.3f}\",\n", " monitor=\"valid_accs\",\n", " mode=\"max\",\n", " save_top_k=1,\n", " save_last=True,\n", " )\n", "\n", "\n", "def my_early_stopping():\n", " # Monitor a metric and stop training when it stops improving.\n", " return L.pytorch.callbacks.EarlyStopping(\n", " monitor=\"valid_accs\",\n", " mode=\"max\",\n", " patience=4,\n", " verbose=True,\n", " )" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def train_model(\n", " model,\n", " experiment_path,\n", " dl_train,\n", " dl_valid,\n", " max_epochs=MAX_EPOCHS,\n", " **trainer_kwargs,\n", "):\n", " callbacks = [\n", " L.pytorch.callbacks.TQDMProgressBar(leave=True),\n", " L.pytorch.callbacks.LearningRateMonitor(),\n", " my_training_checkpoint(),\n", " # Early stopping can be convenient for finding a good\n", " # max_epochs value, but just like ReduceLROnPlateau it\n", " # can make ablation study results less pure/deterministic\n", " # my_early_stopping(),\n", " ]\n", " trainer = L.Trainer(\n", " callbacks=callbacks,\n", " max_epochs=max_epochs,\n", " default_root_dir=experiment_path,\n", " **trainer_kwargs,\n", " )\n", " trainer.fit(model, dl_train, dl_valid)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_model(\n", " MyCIFAR10TrainingModule(),\n", " \"runs/basic_callbacks\",\n", " dl_train,\n", " dl_valid,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Разморозка весов" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Обучение нейросетей требует большого объема данных и много времени. Часто удобно использовать веса, заранее полученные при обучении на некотором большом датасете (ImageNet, OpenImages). В зависимости от размера нашего набора данных, можно тренировать разное количество слоев (начиная с выхода нейросети), а остальные \"заморозить\".\n", "Использовать можно следующую схему:" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "![](https://courses.cv-gml.ru/storage/seminars/nn-training-advanced/freeze_methods.png)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Добавим возможность замораживать часть весов:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_frozen_resnet(num_classes, unfreeze=\"most\", transfer=True):\n", " resnet_model = get_resnet_torchvision(num_classes, transfer)\n", "\n", " unfreeze_params = {\n", " \"last\": -1,\n", " \"most\": -4,\n", " \"full\": 0,\n", " }\n", "\n", " assert unfreeze in unfreeze_params.keys()\n", " first_unfrozen = unfreeze_params[unfreeze]\n", "\n", " for child in list(resnet_model.children())[:first_unfrozen]:\n", " for param in child.parameters():\n", " param.requires_grad = False\n", "\n", " return resnet_model" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class MyCIFAR10TrainingModuleFrozen(MyCIFAR10TrainingModule):\n", " def __init__(self, *, unfreeze=\"most\", **kwargs):\n", " self.unfreeze = unfreeze\n", " super().__init__(**kwargs)\n", "\n", " def get_model(self):\n", " return get_frozen_resnet(\n", " self.num_classes,\n", " self.unfreeze,\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Попробуем обучить 3 нейросети с разным числом \"размороженных\" слоев: только последний слой (last), часть слоев (most), все слои (full)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for unfreeze in [\"last\", \"most\", \"full\"]:\n", " train_model(\n", " MyCIFAR10TrainingModuleFrozen(unfreeze=unfreeze),\n", " f\"runs/unfreeze_{unfreeze}\",\n", " dl_train,\n", " dl_valid,\n", " )" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Warmup" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Процесс дообучения может быть нестабилен в самом начале из-за больших значений функции потерь (из-за специфики самой задачи или из-за случайной инициализации последних слоев в нейросети). В подобных ситуациях попытки сразу начать процесс дообучения, используя большой стартовый learning rate, могут привести к тому, что сеть \"забывает\" часть предобученной информации или даже расходится.\n", "\n", "Warm-up позволяет снизить влияние такого эффекта на начальных этапах. Если при обычном обучении мы бы использовали стартовый learning rate $\\lambda_{start}$, то при использовании warm-up предлагается начинать обучение, используя learning rate $k \\cdot \\lambda_{start}\\ \\ (0 \\lt k \\ll 1)$, и затем в течение $n$ эпох линейно увеличивать его обратно до $\\lambda_{start}$, после чего обучение продолжается так же, как обычно.\n", "\n", "Период warm-up как правило длится одну эпоху." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Скомбинируем Warm-up и Cosine Annealing LR из прошлого семинара:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def plot_lrs(scheduler_cls, show_num_steps, /, **kwargs):\n", " fig, ax = plt.subplots()\n", "\n", " p = torch.zeros(1, requires_grad=True)\n", " dummy_optimizer = torch.optim.Adam([p], lr=0.001)\n", " scheduler = scheduler_cls(dummy_optimizer, **kwargs)\n", " lrs = []\n", " for _ in range(show_num_steps):\n", " lrs.append(scheduler.get_last_lr())\n", " dummy_optimizer.step(), scheduler.step()\n", "\n", " ax.plot(lrs)\n", "\n", " ax.grid()\n", " ax.set_xlabel(\"Step\")\n", " ax.set_ylabel(\"Learning rate\")\n", " kwargs_str = \", \".join(f\"{k}={v!r}\" for k, v in kwargs.items())\n", " ax.set_title(f\"{scheduler_cls.__name__}({kwargs_str})\")\n", "\n", " plt.tight_layout()\n", " plt.show()\n", "\n", "\n", "def warmup_then_cosine_annealing_lr(\n", " optimizer,\n", " start_factor,\n", " total_steps,\n", " warmup_duration,\n", "):\n", " warmup = torch.optim.lr_scheduler.LinearLR(\n", " optimizer,\n", " start_factor=start_factor,\n", " end_factor=1.0,\n", " total_iters=warmup_duration,\n", " )\n", " cos_annealing = torch.optim.lr_scheduler.CosineAnnealingLR(\n", " optimizer,\n", " T_max=total_steps - warmup_duration,\n", " )\n", " warmup_then_cos_anneal = torch.optim.lr_scheduler.SequentialLR(\n", " optimizer,\n", " [warmup, cos_annealing],\n", " milestones=[warmup_duration],\n", " )\n", " return warmup_then_cos_anneal\n", "\n", "\n", "plot_lrs(\n", " warmup_then_cosine_annealing_lr,\n", " 1600,\n", " start_factor=0.1,\n", " total_steps=1600,\n", " warmup_duration=200,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class MyCIFAR10TrainingModuleWarmupCos(MyCIFAR10TrainingModuleFrozen):\n", " def configure_optimizers(self):\n", " optimizer = super().configure_optimizers()\n", "\n", " steps_per_epoch = len(dl_train)\n", " warmup_duration = 1 * steps_per_epoch\n", " total_steps = MAX_EPOCHS * steps_per_epoch\n", "\n", " scheduler = warmup_then_cosine_annealing_lr(\n", " optimizer,\n", " start_factor=0.0001,\n", " total_steps=total_steps,\n", " warmup_duration=warmup_duration,\n", " )\n", "\n", " lr_scheduler = {\n", " \"scheduler\": scheduler,\n", " # The unit of the scheduler's step size, could also be 'epoch'.\n", " # 'epoch' updates the scheduler on epoch end whereas 'step'\n", " # updates it after a optimizer update.\n", " \"interval\": \"step\",\n", " # How many epochs/steps should pass between calls to\n", " # `scheduler.step()`. 1 corresponds to updating the learning\n", " # rate after every `interval` (epoch or step).\n", " \"frequency\": 1,\n", " ## Metric to to monitor for schedulers like `ReduceLROnPlateau`\n", " # \"monitor\": \"valid_accs\",\n", " }\n", " return [optimizer], [lr_scheduler]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_model(\n", " MyCIFAR10TrainingModuleWarmupCos(),\n", " \"runs/warmup_cosine\",\n", " dl_train,\n", " dl_valid,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Обучение адапторов\n", "\n", "В NLP крайне популярно семейство методов дообучения, основанных на добавлении в нейросеть адапторов. Базовым методом из данной категории является метод LoRA (Low-Rank Adaptation), предложенный в [данной статье](https://arxiv.org/pdf/2106.09685).\n", "\n", "Идея подобных методов заключается в том, что вместо прямого обучения параметров нейросети $P$, можно обучать **добавки** к этим параметрам $P + \\Delta$. Дальше, делается предположение, что эти добавки будут имет вид матриц низкого ранга и соответственно будут допускать разложение $\\Delta = A \\cdot B$, где $size(A) + size(B) \\ll size(P)$.\n", "\n", "\n", "\n", "Данное семейство методов в CV тоже используются, но меньше чем в NLP (в основном - при дообучении супер больших трансформерных сетей)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Эффективная загрузка и аугментация данных на GPU" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### NVIDIA DALI\n", "\n", "[NVIDIA DALI](https://developer.nvidia.com/dali) — библиотека, которая позволяет проводить декодирование, аугментирование и прочую обработку данных на GPU. DALI может считывать изображения с диска из большого количества различных [форматов датасетов](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/operations/nvidia.dali.fn.readers.html). Если вашего формата нет в этом списке, то вы можете также использовать загрузчик [`external_source`](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/operations/nvidia.dali.fn.external_source.html) для определения собственной логики загрузки на Python.\n", "\n", "CIFAR10 не очень большой датасет, так что `torchvision.datasets.CIFAR10` на самом деле загружает его целиком в память. В демонстративных целях давайте сохраним изображения CIFAR10 на диск в формате `.jpg` и запишем имена файлов и метки классов в `labels.txt`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def save_CIFAR10_images(ds, root):\n", " os.makedirs(root, exist_ok=True)\n", "\n", " with open(f\"{root}/labels.txt\", \"w\") as f:\n", " for idx, (img, label) in tqdm(enumerate(ds), total=len(ds)):\n", " filename = f\"{idx}.jpg\"\n", " plt.imsave(f\"{root}/{filename}\", img)\n", " f.write(f\"{filename} {label}\\n\")\n", "\n", "\n", "if not os.path.exists(\"cifar10_images\"):\n", " ds_train = CIFAR10(transform=None, train=True)\n", " ds_valid = CIFAR10(transform=None, train=False)\n", "\n", " save_CIFAR10_images(ds_train, \"cifar10_images/train\")\n", " save_CIFAR10_images(ds_valid, \"cifar10_images/valid\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Далее, нам необходимо реализовать описание пайплайна для загрузки и предобработки данных.\n", "\n", "Сохраненные таким образом файлы подходят для загрузки с помощью [`fn.readers.file`](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/operations/nvidia.dali.fn.readers.file.html#nvidia.dali.fn.readers.file). Обратите внимание, что DALI позволяет загружать сырые байты PNG/JPEG изображений прямо на GPU и затем декодировать их уже в памяти графического ускорителя.\n", "\n", "В NVIDIA DALI довольно мало (относительно Albumentations) уже реализованных аугментаций, потому многие операторы нужно собирать из базовых функций." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@pipeline_def(enable_conditionals=True)\n", "def get_pipeline(file_root, is_training):\n", " jpgs, labels = fn.readers.file(\n", " name=\"Reader\",\n", " file_root=file_root,\n", " random_shuffle=is_training,\n", " file_list=f\"{file_root}/labels.txt\",\n", " )\n", " images = fn.decoders.image(jpgs, device=\"mixed\")\n", " labels = labels.gpu()\n", " labels = fn.cast(labels, dtype=types.INT64)\n", "\n", " if is_training:\n", " rotate_limit, p = 15, 0.5\n", " if fn.random.uniform(range=(0, 1)) < p:\n", " angle = fn.random.uniform(range=(-rotate_limit, rotate_limit))\n", " images = fn.rotate(images, angle=angle, keep_size=True)\n", "\n", " r_shift_limit, g_shift_limit, b_shift_limit, p = 15, 15, 15, 0.5\n", " if fn.random.uniform(range=(0, 1)) < p:\n", " shifts = fn.random.uniform(range=(-1, 1), shape=(3,))\n", " shifts = shifts * (r_shift_limit, g_shift_limit, b_shift_limit)\n", " images = images + shifts\n", " images = dali.math.clamp(images, 0, 255)\n", " images = fn.cast(images, dtype=types.UINT8)\n", "\n", " brightness_limit, contrast_limit, p = 0.05, 0.05, 0.5\n", " if fn.random.uniform(range=(0, 1)) < p:\n", " brightness = fn.random.uniform(\n", " range=(1 - brightness_limit, 1 + brightness_limit),\n", " )\n", " contrast = fn.random.uniform(\n", " range=(1 - contrast_limit, 1 + contrast_limit),\n", " )\n", " images = fn.brightness_contrast(\n", " images,\n", " brightness=brightness,\n", " contrast=contrast,\n", " )\n", "\n", " images = (images / 255 - MEAN) / STD\n", "\n", " images = fn.transpose(images, perm=[2, 0, 1])\n", "\n", " return images, labels" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_dali_loader(file_root, is_training, batch_size, num_workers):\n", " pipeline = get_pipeline(\n", " file_root=file_root,\n", " is_training=is_training,\n", " batch_size=batch_size,\n", " num_threads=num_workers,\n", " device_id=0,\n", " )\n", " policy = \"DROP\" if is_training else \"PARTIAL\"\n", " policy = getattr(dali_pytorch.LastBatchPolicy, policy)\n", " loader = dali_pytorch.DALIClassificationIterator(\n", " pipeline,\n", " reader_name=\"Reader\",\n", " last_batch_policy=policy,\n", " )\n", " return loader\n", "\n", "\n", "dl_train_dali = get_dali_loader(\"cifar10_images/train\", True, BATCH_SIZE, NUM_WORKERS)\n", "dl_valid_dali = get_dali_loader(\"cifar10_images/valid\", False, BATCH_SIZE, NUM_WORKERS)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "`DALIClassificationIterator` выдает батчи формата `[{\"data\": ..., \"label\": ...}]`. Для корректной обработки батчей обновим Lightning-модуль:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class MyCIFAR10TrainingModuleDALI(MyCIFAR10TrainingModuleFrozen):\n", " def _step(self, batch, kind):\n", " batch = self.prepare_dali_batch(batch)\n", " return super()._step(batch, kind)\n", "\n", " def prepare_dali_batch(self, batch):\n", " (batch,) = batch\n", " x = batch[\"data\"]\n", " y = batch[\"label\"].squeeze(-1)\n", " return (x, y)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_model(\n", " MyCIFAR10TrainingModuleDALI(),\n", " \"runs/dali\",\n", " dl_train_dali,\n", " dl_valid_dali,\n", " num_sanity_val_steps=0,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Другие библиотеки для ускорения загрузки / обработки данных:\n", "\n", "* [`torchvision.io.decode_jpeg`](https://docs.pytorch.org/vision/main/generated/torchvision.io.decode_jpeg.html) - декодирование `jpeg` файлов на GPU без NVIDIA DALI\n", "\n", "\n", "\n", "* [`GPUDirect`](https://developer.nvidia.com/gpudirect) и [`cuFile`](https://docs.nvidia.com/gpudirect-storage/api-reference-guide/index.html) - более эффективное чтение файлов с диска прямо на GPU (потенциально - вообще в обход CPU)\n", "\n", "\n", "\n", "* [`kornia`](https://kornia.github.io/) - реализация аугментаций как PyTorch операций (дифференцируемые и с поддержкой GPU)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Способы увеличения размера батча" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Multi-device обучение\n", "\n", "В случае, если у вас есть доступ к нескольким графическим ускорителям (или даже нескольким машинам с ускорителями), то достаточно логичной является идея использовать все эти девайсы. Этот подход может позволить вам ускорить сам процесс обучения или даже увеличить размера батча (за счет того что каждый ускоритель/машина имеет свою память).\n", "\n", "Самый простой способ использования нескольких девайсов - это просто одновременный запуск нескольких экспериментов (по одному эксперименту на девайс). К сожалению, данный подход \"ускоряет\" процесс проведения экспериментов, только амортизировано при проведении нескольких экспериментов и вообще не позволяет увеличить размер батча.\n", "\n", "Для того чтобы использовать несколько девайсов или машин в ходе **одного** эксперимента, в PyTorch есть модуль [`torch.distributed`](https://docs.pytorch.org/docs/stable/distributed.html) с реализацией низкоуровневых операций для коммуникаций между девайсами и машинами, а также высокоуровневый класс [`torch.nn.parallel.DistributedDataParallel`](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) (или \"`DDP`\").\n", "\n", "Идея `DDP` является в \"разбиении\" каждого батча между несколькими девайсами. При этом на каждом девайсе будет храниться полная копия всех параметров модели. Следовательно, при обновлении параметров в ходе обучения требуется синхронизировать градиенты или параметры между всеми девайсами.\n", "\n", "Кроме того, добавление `DDP` в обучение требует особой осторожности. В частности - при загрузке данных необходимо учитывать для какого девайса эти данные загружаются. Слои вроде `BatchNorm` или сложные функции потерь, сравнивающие разные примеры внутри каждого батча, могут требовать дополнительной логики для обмена данными между девайсами.\n", "\n", "#### Sharded методы\n", "\n", "Как было упомянуто выше, `DDP` предполагает разбиение батчей / данных между девайсами. При этом сами параметры модели дублируются между девайсами, а вычисления в `forward`/`backward` проводятся одновременно (для разных данных).\n", "\n", "Естественно, это не единственный способ параллелизации вычислений между несколькими девайсами. В NLP крайне популярны методы параллелизации, делающие разбиение самой модели (и соответственно вычислений `forward`/`backward`). Подобные методы называются \"Model Parallel\" или \"Sharded Parallel\".\n", "\n", "Самые актуальные в данный момент методы из этой категории:\n", "\n", "- FSDP (Fully-Sharded Data Parallel)\n", " - [`torch.distributed.fsdp`](https://docs.pytorch.org/docs/stable/fsdp.html)\n", " - [`torch.distributed.fsdp.fully_shard`](https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html)\n", "\n", "\n", "\n", "- Tensor Parallel\n", " - [`torch.distributed.tensor.parallel`](https://docs.pytorch.org/docs/stable/distributed.tensor.parallel.html)\n", " - [`torch.distributed.tensor.DTensor`](https://docs.pytorch.org/docs/stable/distributed.tensor.html#torch.distributed.tensor.DTensor)\n", "\n", "В CV обычно используется простой `DDP`, а \"sharded\" методы используются намного реже. В отличие от NLP моделей, количество параметров в CV моделях обычно увеличивается медленнее чем вычислительная сложность `forward`/`backward` и соответственно \"sharded\" паралеллизм имеет меньше смысла.\n", "\n", "#### В Lightning\n", "\n", "В `lightning`, выше упомянутые методы (и многие другие) называются \"стратегиями\" и могут быть автоматически включены просто правильно оформив обучение и сконфигурировав `Trainer`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Gradient Accumulation\n", "\n", "Пусть в GPU помещается батч размера $N$. Так как выполнение `backward()` не перезаписывает значения градиентов, а аккумулирует их (`+=`), можно выполнять шаг оптимизатора (и последующее обнуление градиентов) после каждого $k$-го батча размера $N$, эмулируя \"аккумулированный\" батч размера $k \\cdot N$.\n", "\n", "При этом обновление градиентного спуска с Gradient Accumulation **не будет** полностью эквивалентно обновлению при размере батча в $k \\cdot N$ — например, из-за слоев типа `BatchNorm`, которые собирают статистики по исходному батчу размера $N$.\n", "\n", "В PyTorch Lightning за Gradient Accumulation отвечает параметр `accumulate_grad_batches` в `Trainer`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_model(\n", " MyCIFAR10TrainingModuleFrozen(),\n", " \"runs/accumulate_grad\",\n", " dl_train,\n", " dl_valid,\n", " accumulate_grad_batches=4,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Также в Lightning можно использовать callback [`L.pytorch.callbacks.GradientAccumulationScheduler`](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.GradientAccumulationScheduler.html) для автоматического изменения количества аккумулированных батчей в ходе обучения." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Gradient Checkpointing\n", "\n", "Во время обучения бóльшую часть выделенной на GPU памяти занимает граф вычислений, который хранит выходные значения всех слоев нейросети. С помощью модуля `torch.utils.checkpoint` можно сохранять выходы только части слоев (например, каждого $k$-го слоя). Тогда во время работы `.backward()` от $i$-го чекпоинта до $(i-1)$-го потребуется выполнить `forward()` и (временно) сохранить выходы от $(i - 1)$-го чекпоинта до $i$-го.\n", "\n", "Как итог, чекпоинтинг градиентов может уменьшить потребление памяти ценой выполнения дополнительных вычислений.\n", "\n", "В отличие от NLP, данный прием не так часто используется на практике в CV.\n", "\n", "Больше про этот прием можно узнать в [документации PyTorch](https://pytorch.org/docs/stable/checkpoint.html)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### (Automatic) Mixed Precision\n", "\n", "Данный метод был изначально предложен в статье [Mixed Precision Training](https://arxiv.org/abs/1710.03740) от NVIDIA.\n", "\n", "По умолчанию для хранения весов и выполнения операций используется тип `float32`. Использование типа `float16` может потенциально ускорить вычисления и снизить требования к объему и скорости памяти. Однако просто конвертировать все веса и входы в `float16` и проводить все вычисления в этом типе скорее всего не получится, т.к. точность и диапазон значений у типа `float16` значительно меньше, чем у `float32`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Из-за ограниченного диапазона и точности типа `float16` всплывают **три** проблемы, исправлением которых и занимается метод Automatic Mixed Precision.\n", "\n", "**Первая проблема**: При попытке использования обычных градиентных оптимизаторов для обновления `float16` параметров может получиться, что шаг градиентного спуска $\\lambda\\nabla L$ будет настолько мал по сравнению c $P$, что с учетом округления $P - \\lambda\\nabla L = P$ (то есть параметры не обновляются или обновляются плохо).\n", "\n", "Основная идея Mixed Precision Training, исправляющая данную проблему, заключается в том, что при обучении можно хранить \"истинные\" (master) значения весов в `float32`, при вычислении прямого и обратного прохода временно конвертировать их в `float16`, а затем полученные `float16` градиенты обратно конвертировать в `float32` для обновления параметров." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Вторая проблема**: Даже если выполнять обновление параметров, используя тип `float32`, в прямом и обратном проходе нейросети все равно останутся некоторые операции, для которых потеря точности `float16` будет значительной.\n", "\n", "В первую очередь это операции, которые в ходе своего вычисления выполняют редукцию (например, матричное умножение и свертка). В подобных операциях во время многократного суммирования будет накапливаться ошибка. К счастью, в Tensor Cores современных графических ускорителей реализуются специальные операции, позволяющие проводить основную часть вычислений в типе `float16`, а затем аккумулировать результат уже в типе `float32`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Кроме того, потери точности или переполнения могут возникать в обычных агрегирующих операциях (например, `.sum()` и `.mean()`), а также при использовании примитивных операций, которые сами по себе менее стабильны (например, `exp`, `log` и `1/x`). Подобные операции обычно встречаются в некоторых функциях активаций, в функциях потерь и в различных видах нормализации.\n", "\n", "Для таких операций предлагается конвертировать `float16` входы/параметры в `float32`, проводить вычисления в расширенном типе, а затем возможно конвертировать их обратно в `float16`. В `PyTorch` уже реализован функционал, автоматически определяющий, в каком типе лучше всего проводить вычисления для большинства распространенных нейросетевых операций. Для этого вместо ручного приведения к нужным типам необходимо просто обернуть интересующие нас вызовы в контекстный менеджер `torch.autocast`.\n", "\n", "В ситуациях, когда вы хотите самостоятельно реализовать какие-то сложные операции и более явно контролировать для них типы промежуточных вычислений, вы можете использовать декораторы `torch.amp.custom_fwd`/`custom_bwd` в классах `torch.autograd.Function` и контекстный менеджер `torch.autocast(enable=True/False)` для обычных вычислений. См. раздел [Autocast and Custom Autograd Functions](https://pytorch.org/docs/stable/notes/amp_examples.html#autocast-and-custom-autograd-functions) в документации Automatic Mixed Precision." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class MyCIFAR10TrainingModuleAutocast(MyCIFAR10TrainingModuleFrozen):\n", " def __init__(self, *, autocast_dtype=torch.float32, **kwargs):\n", " super().__init__(**kwargs)\n", " self.autocast_dtype = autocast_dtype\n", "\n", " def training_step(self, batch):\n", " x, y = batch\n", "\n", " with torch.autocast(\n", " device_type=self.device.type,\n", " dtype=self.autocast_dtype,\n", " ):\n", " # Dynamically mix fp32 and fp16\n", " # during the forward computation\n", " p = self.model(x)\n", " loss = self.loss_fn(p, y)\n", "\n", " accs = self.accuracy(p.argmax(axis=-1), y)\n", " return self._log_metrics(loss, accs, \"train\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Третья проблема**: Кроме проблем с погрешностями при вычислении градиентов может оказаться, что даже точно вычисленные градиенты имеют распределение, выходящее за границы диапазона чисел, представимых в типе `float16`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ " " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Обнуления/переполнения градиентов в `float16` можно избежать, если во время обратного распространения ошибки вычислять не истинные значения градиентов, а умноженные на какую-то константу $S$, а затем во время шага оптимизатора умножать уже сконвертированные в `float32` градиенты на $1/S$. Таким образом, `float16` градиенты вычисляются с использованием смещенного диапазона значений." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "При правильном выборе множителя $S$ подобное обучение может сходиться так же хорошо, как с использованием `float32`, но работая быстрее и используя меньше памяти." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Тем не менее, не очень удобно подбирать правильное значение $S$ вручную. Для автоматизации процедуры поиска данного множителя было бы удобно уметь эффективно определять факт обнуления или переполнения в ходе вычислений с плавающей точкой. На практике числа с плавающей точкой IEEE 754 не предусматривают функционал для эффективного определения обнуления.\n", "\n", "С другой стороны, для автоматического определения переполнения значений достаточно после обратного прохода проверить, есть ли среди вычисленных градиентов $\\pm\\infty$ или $\\texttt{NaN}$. Тогда, используя данную информацию, можно реализовать следующий алгоритм:\n", "\n", "1. Выполняем прямой проход используя описанные выше mixed precision вычисления.\n", "\n", "2. Умножаем `loss` на множитель $S$ и выполняем обратный проход. В итоге получаем `float16` градиенты.\n", "\n", "3. Если среди вычисленных градиентов есть $\\pm\\infty$ или $\\texttt{NaN}$\n", " - Cчитаем, что значение $S$ было слишком большим — делим $S$ на 2, выкидываем посчитанные градиенты, пропускаем текущий батч и пробуем ещё раз.\n", "

\n", "\n", "4. Если в градиентах не было $\\pm\\infty$ и $\\texttt{NaN}$\n", " - Конвертируем градиенты в `float32`, умножаем их на $1/S$ и выполняем шаг оптимизации как при обычном обучении.\n", "

\n", "\n", "5. Если $\\pm\\infty$ и $\\texttt{NaN}$ не было уже много шагов подряд\n", " - Значит значение $S$ может быть слишком маленьким — пробуем умножить $S$ на 2 и смотрим, получатся ли $\\pm\\infty$ или $\\texttt{NaN}$." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "В `PyTorch` данный функционал реализован в классе `torch.amp.GradScaler`:\n", "\n", "```python\n", "scaler = torch.amp.GradScaler()\n", "\n", "for epoch in epochs:\n", " for x, y in data:\n", " optimizer.zero_grad()\n", " with torch.autocast(...):\n", " p = model(x)\n", " loss = loss_fn(p, y)\n", "\n", " # Scales loss. Calls backward() on scaled loss to create scaled gradients.\n", " scaler.scale(loss).backward()\n", "\n", " # scaler.step() first unscales gradients of the optimizer's params.\n", " # If gradients don't contain infs/NaNs, optimizer.step() is then called,\n", " # otherwise, optimizer.step() is skipped.\n", " scaler.step(optimizer)\n", "\n", " # Updates the scale for next iteration.\n", " scaler.update()\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "В PyTorch Lightning типы для вычислений можно задать параметром `precision` в `Trainer`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_model(\n", " MyCIFAR10TrainingModuleAutocast(autocast_dtype=torch.float16),\n", " \"runs/amp_lightning_fp16\",\n", " dl_train,\n", " dl_valid,\n", " precision=\"16-mixed\",\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "При использовании Automatic Mixed Precision вместе с Gradient Accumulation, градиенты накапливаются в scaled-режиме (**до** приведения к исходному масштабу). [Пример в PyTorch](https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Тип bfloat16\n", "\n", "Некоторые современные ускорители аппаратно поддерживают тип `bfloat16` (изначально [разработанный](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus) Google). В `bfloat16` на экспоненту отводится больше бит, чем в `float16`, расширяя диапазон значений (ценой точности). Таким образом, при использовании `bfloat16` можно не делать масштабирование градиентов (не использовать `GradScaler`).\n", "\n", "Однако стоит понимать, что эффективные операции для работы с данным типом доступны не на всех девайсах. В случае отсутствия аппаратной поддержки, вычисления в данном типе могут работать даже медленнее чем `float32`." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "train_model(\n", " MyCIFAR10TrainingModuleAutocast(autocast_dtype=torch.bfloat16),\n", " \"runs/amp_lightning_bf16\",\n", " dl_train,\n", " dl_valid,\n", " precision=\"bf16-mixed\",\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Больше про Automatic Mixed Precision можно почитать\n", "- NVIDIA\n", " - [Video Series: Mixed-Precision Training Techniques Using Tensor Cores for Deep Learning](https://developer.nvidia.com/blog/video-mixed-precision-techniques-tensor-cores-deep-learning/)\n", " - [Train With Mixed Precision](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html)\n", "

\n", "- PyTorch\n", " - [`amp` Package Documentation](https://pytorch.org/docs/stable/amp.html)\n", " - [Recipe](https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html)\n", " - [AMP Examples](https://pytorch.org/docs/stable/notes/amp_examples.html)\n", "

\n", "- Lightning\n", " - [N-Bit Precision](https://lightning.ai/docs/pytorch/stable/common/precision_intermediate.html)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---" ] } ], "metadata": {}, "nbformat": 4, "nbformat_minor": 4 }