В этом уроке мы научимся использовать nn.parallel.DistributedDataParallel для обучения наших моделей на нескольких графических процессорах. Мы возьмем минимальный пример обучения классификатора изображений и посмотрим, как мы можем ускорить обучение.

Начнем с импорта.

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import time

Мы будем использовать CIFAR10 во всех наших экспериментах с размером пакета 256.

def create_data_loader_cifar10():
   transform = transforms.Compose(
       (
       transforms.RandomCrop(32),
       transforms.RandomHorizontalFlip(),
       transforms.ToTensor(),
       transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))))
   batch_size = 256
   trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                           download=True, transform=transform)
   trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                           shuffle=True, num_workers=10, pin_memory=True)
   testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
   testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                           shuffle=False, num_workers=10)
   return trainloader, testloader

Сначала мы будем обучать модель на одном графическом процессоре Nvidia A100 в течение 1 эпохи. Стандартный pytorch материал здесь, ничего нового. Учебник основан на официальный учебник из документов Pytorch.

def train(net, trainloader):
   print("Start training...")
   criterion = nn.CrossEntropyLoss()
   optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
   epochs = 1
   num_of_batches = len(trainloader)
   for epoch in range(epochs):  
       running_loss = 0.0
       for i, data in enumerate(trainloader, 0):
           inputs, labels = data
           images, labels = inputs.cuda(), labels.cuda()
           optimizer.zero_grad()
           outputs = net(images)
           loss = criterion(outputs, labels)
           loss.backward()
           optimizer.step()
           running_loss += loss.item()
       print(f'(Epoch {epoch + 1}/{epochs}) loss: {running_loss / num_of_batches:.3f}')
   print('Finished Training')

test аналогично определяется функция. Основной скрипт просто соберет все вместе:

if __name__ == '__main__':
   start = time.time()  
   PATH = './cifar_net.pth'
   trainloader, testloader = create_data_loader_cifar10()
   net = torchvision.models.resnet50(False).cuda()
   start_train = time.time()
   train(net, trainloader)
   end_train = time.time()
   torch.save(net.state_dict(), PATH)
   test(net, PATH, testloader)
   end = time.time()
   seconds = (end - start)
   seconds_train = (end_train - start_train)
   print(f"Total elapsed time: {seconds:.2f} seconds, \
       Train 1 epoch {seconds_train:.2f} seconds")

Мы используем resnet50 для измерения производительности сети приличного размера.

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

$ python -m train_1gpu

Accuracy of the network on the 10000 test images: 27 %
Total elapsed time: 69.03 seconds,      Train 1 epoch 13.08 seconds

Ладно, пора заняться оптимизацией.

Код доступен на Гитхаб. Если вы планируете закрепить свои знания о Pytorch, мы настоятельно рекомендуем две замечательные книги: Глубокое обучение с PyTorch от Manning Publications и Машинное обучение с PyTorch и Scikit-Learn Себастьян Рашка. Вы всегда можете использовать код скидки 35% блейсаммер21 для всей продукции Manning’s.

torch.nn.DataParallel: без боли, без выигрыша

DataParallel является однопроцессным, многопоточным и работает только на одной машине. Для каждого графического процессора мы используем одну и ту же модель для прямого прохода. Мы распределяем данные по графическим процессорам и выполняем прямые проходы в каждом из них. По сути, происходит то, что размер партии делится на количество рабочих.

В этом варианте использования эта функциональность не давала никакой выгоды. Это потому, что система, которую я использую, имеет узкие места в процессоре и жестком диске. Другие машины с очень быстрыми дисками и процессором, но с трудом справляющиеся со скоростью графического процессора (узким местом графического процессора), могут извлечь выгоду из этой функции.

На практике единственное изменение, которое вам нужно сделать в коде, это следующее:

net = torchvision.models.resnet50(False)
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    net = nn.DataParallel(net)

Когда используешь nn.DataParallelразмер партии должен быть делимый по количеству графических процессоров.

nn.DataParallel разделяет пакет и обрабатывает его независимо на всех доступных графических процессорах. При каждом прямом проходе модуль реплицируется на каждом графическом процессоре, что приводит к значительным накладным расходам. Каждая реплика обрабатывает часть пакета (batch_size/gpus). Во время обратного прохода градиенты от каждой реплики суммируются в исходный модуль.

Дополнительная информация в нашей предыдущей статье о параллелизме данных и модели.

Хорошей практикой при использовании нескольких графических процессоров является заранее определить графические процессоры, которые будет использовать ваш скрипт:

import os
os.environ('CUDA_VISIBLE_DEVICES') = "0,1"

Это должно быть СДЕЛАНО до любой другой импорт, связанный с CUDA.

Даже из Pytorch документация очевидно, что это очень плохая стратегия:

Рекомендуется использовать nn.DistributedDataParallelвместо этого класса проводить обучение с несколькими GPU, даже если есть только один узел.

Причина в том, что DistributedDataParallel использует один процесс на одного работника (GPU), а DataParallel инкапсулирует всю передачу данных в один процесс.

Согласно документам, данные могут находиться на любом устройстве до того, как они будут переданы в модель.

В моем эксперименте DataParallel был помедленнее чем обучение на одном графическом процессоре. Даже с 4 GPU. После увеличения количества рабочих я уменьшил время, но все равно хуже, чем одиночный GPU. Я измеряю и сообщаю время, необходимое для обучения модели для одной эпохи, то есть 50 000 изображений 32×32.

Последнее замечание: чтобы сравнить производительность с одним GPU, я умножил размер пакета на количество рабочих, то есть 4 для 4 GPU. В противном случае это более чем в 2 раза медленнее.

Это подводит нас к жесткой теме распределенных параллельных данных.

Код доступен на Гитхаб. Вы всегда можете поддержать нашу работу, поделившись в социальных сетях, сделав пожертвование и купив нашу книгу и электронный курс.

Распределенная параллельная передача данных Pytorch

Распределенная параллельная обработка данных является многопроцессорной и работает как для обучения на одной, так и на нескольких машинах. в питорке, nn.parallel.DistributedDataParallel распараллеливает модуль, разделяя ввод между указанными устройствами. Этот модуль также подходит для обучения работе с несколькими узлами и несколькими графическими процессорами. Здесь я экспериментировал только с одним узлом (1 машина с 4 графическими процессорами).

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

Модуль реплицируется на каждую машину и каждое устройство. Во время прямого прохода каждый рабочий процесс (GPU) обрабатывает данные и локально вычисляет собственный градиент. Во время обратного прохода градиенты от каждого узла усредняются. Наконец, каждый рабочий выполняет обновление параметра и отправляет всем остальным узлам вычисленное обновление параметра.

Модуль выполняет все-уменьшить шаг на градиенты и предполагает, что они будут модифицироваться оптимизатором во всех процессах одинаково.

Ниже приведены рекомендации по преобразованию вашего сценария с одним GPU в обучение с несколькими GPU.

Шаг 1: Инициализируйте процессы распределенного обучения

def init_distributed():
    dist_url = "env://" 
    rank = int(os.environ("RANK"))
    world_size = int(os.environ('WORLD_SIZE'))
    local_rank = int(os.environ('LOCAL_RANK'))
    dist.init_process_group(
            backend="nccl",
            init_method=dist_url,
            world_size=world_size,
            rank=rank)
    torch.cuda.set_device(local_rank)
    dist.barrier()

Эта инициализация работает, когда мы запускаем наш скрипт с torch.distributed.launch (Питорч 1.7 и 1.8) или torch.run (Pytorch 1.9+) с каждого узла (здесь 1).

Шаг 2: Оберните модель с помощью DDP

net = torchvision.models.resnet50(False).cuda()
net = nn.SyncBatchNorm.convert_sync_batchnorm(net)
local_rank = int(os.environ('LOCAL_RANK'))
net = nn.parallel.DistributedDataParallel(net, device_ids=(local_rank))

Если каждый процесс имеет правильный локальный ранг, tensor.cuda() или model.cuda() может быть вызван правильно во всем скрипте.

Шаг 3: Используйте DistributedSampler в вашем DataLoader

import torch
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
import torch.nn as nn
def create_data_loader_cifar10():
    transform = transforms.Compose(
        (
        transforms.RandomCrop(32),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))))
    batch_size = 256
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)                                  
    train_sampler = DistributedSampler(dataset=trainset, shuffle=True)                                                  
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                            sampler=train_sampler, num_workers=10, pin_memory=True)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                        download=True, transform=transform)
    test_sampler =DistributedSampler(dataset=testset, shuffle=True)                                         
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                            shuffle=False, sampler=test_sampler, num_workers=10)
    return trainloader, testloader

В распределенном режиме вызов data_loader.sampler.set_epoch() метод в начале каждой эпохи до создание DataLoader итератор необходим для правильной работы перетасовки в нескольких эпохах. В противном случае всегда будет использоваться один и тот же порядок.

def train(net, trainloader):
    print("Start training...")
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    epochs = 1
    num_of_batches = len(trainloader)
    for epoch in range(epochs): 
        trainloader.sampler.set_epoch(epoch)

В более общем виде:

for epoch in range(epochs):
        data_loader.sampler.set_epoch(epoch)
        train_one_epoch(...)

Передовой опыт для DDP

Любые методы, загружающие данные, должны быть изолированы от главного процесса. Любые методы, выполняющие файловый ввод-вывод, должны быть изолированы от главного процесса.

import torch.distributed as dist
import torch 
def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True
def save_on_master(*args, **kwargs):
    if is_main_process():
        torch.save(*args, **kwargs)
def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()
def is_main_process():
    return get_rank() == 0

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

if is_main_process():

Запустить скрипт с помощью torch.distributed.launch или torch.run

$ python -m torch.distributed.launch --nproc_per_node=4 main_script.py

Ошибки будут. Обязательно отключите любой нежелательный распределенный процесс обучения:

$ kill $(ps aux | grep main_script.py | grep -v grep | awk '{print $2}')

Заменять main_script.py с именем вашего скрипта. Еще один более простой вариант $ kill -9 PID. В противном случае вы можете перейти к более сложным вещам, например, убить все процессы, связанные с графическим процессором CUDA, если они не показаны в nvidia-smi

lsof /dev/nvidia* | awk '{print $2}' | xargs -I {} kill {}

Это только в том случае, если вы не можете найти PID процесса, запущенного в графическом процессоре.

Очень хорошая книга по распределенному обучению Распределенное машинное обучение с помощью Python: ускорение обучения и обслуживания моделей с помощью распределенных систем Гуаньхуа Ван.

Обучение смешанной точности в Pytorch

Смешанная точность объединяет вычисления с плавающей запятой (FP) 16 и FP 32 на разных этапах обучения. Тренировка FP16 также известна как тренировка с половинной точностью, которая имеет более низкую производительность. Автоматическая смешанная точность — это буквально лучшее из обоих миров: сокращенное время обучения с производительностью, сравнимой с FP32.

В Тренировка смешанной точности, все вычислительные операции (прямой проход, обратный проход, весовые градиенты) см. в литой версии FP16. Для этого необходима копия веса FP32, а также вычисление потерь в FP32 после прямого прохода в FP16, чтобы избежать переполнения и недополнения. Градиенты весов возвращаются к FP32 для обновления весов модели. Кроме того, потери в FP32 масштабируются, чтобы избежать потери значимости градиента перед переходом к FP16 для выполнения обратного прохода. В качестве компенсации веса FP32 будут уменьшены на один и тот же скаляр перед обновлением веса.

Вот изменения в функции поезда:

fp16_scaler = torch.cuda.amp.GradScaler(enabled=True)
    for epoch in range(epochs):  
        trainloader.sampler.set_epoch(epoch)
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            images, labels = inputs.cuda(), labels.cuda() 
            optimizer.zero_grad()
            with torch.cuda.amp.autocast():
                outputs = net(images)
                loss = criterion(outputs, labels)
            fp16_scaler.scale(loss).backward()
            fp16_scaler.step(optimizer)
            fp16_scaler.update()

Итоги и подведение итогов

В утопическом параллельном мире N рабочих дали бы ускорение N. Здесь вы видите, что вам нужно 4 графических процессора в режиме DistributedDataParallel, чтобы получить ускорение в 2 раза. Обучение смешанной точности обычно обеспечивает значительное ускорение, но графический процессор A100 и другие архитектуры графических процессоров на базе Ampere имеют ограниченный прирост (насколько я читал В сети).

В приведенных ниже результатах указано время в секундах для 1 эпохи на CIFAR10 с resnet50 (размер пакета 256, память графического процессора NVidia A100 40 ГБ):

Очень важное замечание здесь заключается в том, что РаспределенныйДанныеПараллельный использует эффективный размер пакета 4 * 256 = 1024, поэтому он делает меньше обновлений моделей. Вот почему я считаю, что точность проверки гораздо ниже (14% по сравнению с 27% в базовой версии).

Код доступен на Гитхаб если вы хотите поиграть. Результаты будут зависеть от вашего оборудования. Всегда бывает так, что я что-то упустил в своих экспериментах. Если вы обнаружите недостаток, пожалуйста, сообщите мне об этом на нашем дискорд сервер.

Эти выводы дадут вам хороший старт для обучения ваших моделей. Надеюсь, вы найдете их полезными. Поддерживает нас, делясь социальными сетями, делая пожертвования, покупая нашу книгу или электронный курс. Ваша помощь поможет нам производить больше бесплатного контента и доступного контента с искусственным интеллектом. Как всегда, спасибо за интерес к нашему блогу.

Книга «Глубокое обучение в производстве» 📖

Узнайте, как создавать, обучать, развертывать, масштабировать и поддерживать модели глубокого обучения. Изучите инфраструктуру машинного обучения и MLOps на практических примерах.

Узнать больше

* Раскрытие информации: обратите внимание, что некоторые из приведенных выше ссылок могут быть партнерскими ссылками, и мы без дополнительных затрат для вас получим комиссию, если вы решите совершить покупку после перехода по ссылке.