Обучение моделей с помощью Accelerator

Эффективное обучение моделей ИИ с PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer и Accelerator

График сравнения простоты использования и возможности настройки для Accelerator и Trainer.

Пользовательские циклы обучения

Trainer: без пользовательских циклов обучения
Некоторые задачи генеративного ИИ требуют двух сетей

Изображение генеративно-состязательной сети.

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

Trainer и Accelerator

График сравнения простоты использования и возможности настройки для Accelerator и Trainer.

Изменение базового цикла обучения

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

Обнуление градиентов
Перемещение данных на устройство: .to(device)
Прямой проход
Вычисление кросс-энтропийных потерь
Вычисление градиентов (обратный проход)
Обновление параметров модели и скорости обучения

Создание объекта Accelerator

Accelerator предоставляет интерфейс для распределённого обучения

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement (bool, по умолчанию True): автоматическое размещение на устройстве

Определение модели и оптимизатора

Загрузка предобученной модели

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

Оптимизация параметров модели с помощью Adam

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

Определение планировщика

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer (obj): оптимизатор PyTorch, например Adam
num_warmup_steps (int): шаги линейного увеличения lr, задаётся как int(num_training_steps * 0.1)
num_training_steps (int): общее число шагов обучения, задаётся как len(train_dataloader) * num_epochs

Подготовка модели к эффективному обучению

Метод prepare управляет размещением на устройстве

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

Построение цикла обучения с Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

Обнуление градиентов
Ранее данные перемещались на устройство

Построение цикла обучения с Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Обнуление градиентов
Ранее данные перемещались на устройство
Удаление строк с ручным перемещением данных

Построение цикла обучения с Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Обнуление градиентов
Ранее данные перемещались на устройство
Удаление строк с ручным перемещением данных
Прямой проход
Вычисление кросс-энтропийных потерь и градиентов

Построение цикла обучения с Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Обнуление градиентов
Ранее данные перемещались на устройство
Удаление строк с ручным перемещением данных
Прямой проход
Вычисление кросс-энтропийных потерь и градиентов
Замена loss.backward на accelerator
Обновление параметров модели и скорости обучения

Сводка изменений

До Accelerator

Данные перемещались на устройство вручную
- inputs.to(device)
- targets.to(device)
Вычисление градиентов через loss.backward()

После Accelerator

Автоматическое размещение данных и параллелизм
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Синхронизация градиентов через accelerator.backward(loss)
Настраиваемый цикл обучения
Удобный, аппаратно-независимый, масштабируемый и поддерживаемый

Давайте попрактикуемся!

Эффективное обучение моделей ИИ с PyTorch