Accelerate documentation

Запуск распределённого обучения из Jupyter Notebooks

Accelerate

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Запуск распределённого обучения из Jupyter Notebooks

В этом руководстве показано, как дообучить модель компьютерного зрения с помощью 🤗 Accelerate из Jupyter Notebook на распределённой системе. Вы также узнаете, как настроить несколько требований, необходимых для того, чтобы убедиться, что ваша среда настроена правильно, данные подготовлены надлежащим образом, и как запустить обучение.

Это руководство также доступно в формате Jupyter Notebook здесь

Настройка окружения

Прежде чем можно будет запустить обучение, в системе должен существовать файл конфигурации Accelerate. Обычно это делается с помощью следующей команды в терминале с последующим ответом на вопросы:

accelerate config

Однако, если общие значения по умолчанию подходят, и вы не работаете на TPU, в Accelerate есть утилита для быстрого записи конфигурации устройства в файл конфигурации через utils.write_basic_config().

Следующий код перезапустит Jupyter после записи конфигурации, так как для её выполнения был вызван CUDA runtime или XPU runtime.

CUDA и XPU нельзя инициализировать более одного раза на системе с несколькими устройствами. Отладка в ноутбуке с вызовами CUDA/XPU допустима, но для окончательного обучения потребуется выполнить полную очистку и перезапуск.

import os
from accelerate.utils import write_basic_config

write_basic_config()  # Write a config file
os._exit(00)  # Restart the notebook

Подготовка набора данных и модели

Далее вам следует подготовить набор данных. Как уже упоминалось, при подготовке DataLoaders и модели необходимо проявлять особую тщательность, чтобы убедиться, что ничто не размещается на каком-либо GPU.

Если вы делаете это, рекомендуется вынести соответствующий код в отдельную функцию и вызывать её из интерфейса запуска ноутбука, который будет рассмотрен ниже.

Убедитесь, что набор данных загружен в соответствии с инструкциями здесь

import os, re, torch, PIL
import numpy as np

from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator
from accelerate.utils import set_seed
from timm import create_model

Сначала вам нужно создать функцию для извлечения имени класса на основе имени файла:

import os

data_dir = "../../images"
fnames = os.listdir(data_dir)
fname = fnames[0]
print(fname)

beagle_32.jpg

В данном случае метка — beagle. С помощью регулярных выражений можно извлечь метку из имени файла:

import re


def extract_label(fname):
    stem = fname.split(os.path.sep)[-1]
    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]

extract_label(fname)

И вы можете увидеть, что он корректно вернул правильное имя для нашего файла:

"beagle"

Далее следует создать класс Dataset для получения изображения и метки:

class PetsDataset(Dataset):
    def __init__(self, file_names, image_transform=None, label_to_id=None):
        self.file_names = file_names
        self.image_transform = image_transform
        self.label_to_id = label_to_id

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

    def __getitem__(self, idx):
        fname = self.file_names[idx]
        raw_image = PIL.Image.open(fname)
        image = raw_image.convert("RGB")
        if self.image_transform is not None:
            image = self.image_transform(image)
        label = extract_label(fname)
        if self.label_to_id is not None:
            label = self.label_to_id[label]
        return {"image": image, "label": label}

Теперь создадим набор данных. Вне функции обучения можно найти и объявить все имена файлов и метки и использовать их как ссылки внутри запущенной функции:

fnames = [os.path.join("../../images", fname) for fname in fnames if fname.endswith(".jpg")]

Далее соберите все метки:

all_labels = [extract_label(fname) for fname in fnames]
id_to_label = list(set(all_labels))
id_to_label.sort()
label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}

Далее вам следует создать функцию get_dataloaders, которая будет возвращать созданные вами загрузчики данных. Как упоминалось ранее, если данные автоматически отправляются на устройство GPU или TPU при создании ваших DataLoaders, они должны быть созданы с помощью этого метода.

def get_dataloaders(batch_size: int = 64):
    "Builds a set of dataloaders with a batch_size"
    random_perm = np.random.permutation(len(fnames))
    cut = int(0.8 * len(fnames))
    train_split = random_perm[:cut]
    eval_split = random_perm[cut:]

    # For training a simple RandomResizedCrop will be used
    train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])
    train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)

    # For evaluation a deterministic Resize will be used
    eval_tfm = Compose([Resize((224, 224)), ToTensor()])
    eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)

    # Instantiate dataloaders
    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)
    return train_dataloader, eval_dataloader

Наконец, импортируйте планировщик, который будет использоваться позже:

from torch.optim.lr_scheduler import CosineAnnealingLR

Создание функции обучения

Теперь можно создать цикл обучения. notebook_launcher() работает, передавая функцию, которая будет вызвана на распределённой системе.

Вот базовый цикл обучения для задачи классификации животных:

Код разделён на части, чтобы дать пояснения к каждому разделу. Полная версия, которую можно скопировать и вставить, будет представлена в конце

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    accelerator = Accelerator(mixed_precision=mixed_precision)

Сначала следует задать seed и создать объект Accelerator как можно раньше в цикле обучения.

Если обучение выполняется на TPU, цикл обучения должен принимать модель в качестве параметра, а её экземпляр должен создаваться вне функции цикла обучения. Подробнее о причинах — в разделе рекомендации по работе с TPU

Далее вам следует создать загрузчики данных и инициализировать модель:

    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

Здесь вы строите модель так, чтобы seed также контролировал инициализацию новых весов

Поскольку в этом примере выполняется перенос обучения, энкодер модели изначально заморожен, чтобы в начале обучения можно было обучать только головную часть модели:

    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

Нормализация батчей изображений немного ускорит обучение:

    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

Чтобы сделать эти константы доступными на активном устройстве, установите его в качестве устройства объекта Accelerator:

    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

Далее создайте остальные классы PyTorch, используемые для обучения:

    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

Перед передачей всего в prepare().

Особого порядка запоминать не нужно: достаточно распаковать объекты в том же порядке, в котором вы передали их методу prepare.

    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

Теперь обучите модель:

    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

Цикл оценки будет немного отличаться от цикла обучения. Количество элементов, а также общая точность каждого батча будут добавлены к двум константам:

        model.eval()
        accurate = 0
        num_elems = 0

Далее идёт остальная часть стандартного цикла PyTorch:

        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)

И наконец, последняя существенная разница.

При выполнении распределённой оценки предсказания и метки необходимо передать через gather(), чтобы все данные стали доступны на текущем устройстве и можно было получить корректно рассчитанную метрику:

            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

Теперь вам нужно вычислить фактическую метрику для этой задачи, и вы можете вывести её в основном процессе с помощью print():

        eval_metric = accurate.item() / num_elems
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")

Полная версия этого цикла обучения приведена ниже:

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    # Initialize accelerator
    accelerator = Accelerator(mixed_precision=mixed_precision)
    # Build dataloaders
    train_dataloader, eval_dataloader = get_dataloaders(batch_size)

    # Instantiate the model (you build the model here so that the seed also controls new weight initializations)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

    # Freeze the base model
    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

    # You can normalize the batches of images to be a bit faster
    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

    # To make these constants available on the active device, set it to the accelerator device
    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

    # Instantiate the optimizer
    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)

    # Instantiate the learning rate scheduler
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

    # Prepare everything
    # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    # Now you train the model
    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

        model.eval()
        accurate = 0
        num_elems = 0
        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)
            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

        eval_metric = accurate.item() / num_elems
        # Use accelerator.print to print only on the main process.
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")

Использование notebook_launcher

Осталось только использовать notebook_launcher().

Вы передаёте функцию, аргументы (в виде кортежа) и количество процессов для обучения. (Подробнее в документации)

from accelerate import notebook_launcher

args = ("fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=2)

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

Для окружения, содержащего 2 узла (компьютера) с 8 GPU на каждом и основным компьютером с IP-адресом “172.31.43.8”, это будет выглядеть так:

notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=0, num_nodes=2, num_processes=8)

А во второй сессии Jupyter на другой машине:

Обратите внимание, как изменилось значение node_rank

notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=1, num_nodes=2, num_processes=8)

При запуске на TPU это будет выглядеть так:

model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

args = (model, "fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=8)

Для запуска процесса обучения с эластичностью, обеспечивающей отказоустойчивость, можно использовать функцию elastic_launch, предоставляемую PyTorch. Для этого необходимо задать дополнительные параметры, такие как rdzv_backend и max_restarts. Ниже приведен пример использования notebook_launcher с эластичными возможностями:

notebook_launcher(
    training_loop,
    args,
    num_processes=2,
    max_restarts=3
)

Во время выполнения будет выводиться прогресс, а также будет указано, на сколько устройствах была запущена задача. Данный учебник был выполнен на двух GPU:

Launching training on 2 GPUs.
epoch 0: 88.12
epoch 1: 91.73
epoch 2: 92.58
epoch 3: 93.90
epoch 4: 94.71

И это всё!

Обратите внимание, что notebook_launcher() игнорирует файл конфигурации Accelerate. Чтобы запустить на основе конфигурации, используйте:

accelerate launch

Отладка

Частая проблема при запуске notebook_launcher — ошибка о том, что CUDA/XPU уже была инициализирована. Обычно это вызвано импортом или предыдущим кодом в ноутбуке, который обращается к подбиблиотекам PyTorch torch.cuda или torch.xpu. Чтобы сузить поиск причины, вы можете запустить notebook_launcher с переменной окружения ACCELERATE_DEBUG_MODE=yes: при создании процесса будет выполнена дополнительная проверка на то, что обычный процесс может быть создан и использовать CUDA/XPU без проблем. (Ваш код для CUDA/XPU по-прежнему можно будет выполнить после этого).

Заключение

В этом ноутбуке показано, как выполнять распределённое обучение изнутри Jupyter Notebook. Ключевые моменты, которые стоит запомнить:

Убедитесь, что весь код, использующий CUDA/XPU (или импортирующий CUDA/XPU), сохранён для функции, переданной в notebook_launcher()
Установите num_processes равным количеству устройств, используемых для обучения (например, числу GPU, XPU, CPU, TPU и т. д.)
Если вы используете TPU, объявляйте модель вне функции цикла обучения

Update on GitHub

←Запуск скриптов Accelerate Начните здесь!→