Анализ и оптимизация производительности модели PyTorch

Как определить и уменьшить нагрузку на ЦП на этапе обучения с помощью PyTorch Profiler и TensorBoard

Это вторая часть серии постов на тему анализа и оптимизации модели PyTorch, работающей на GPU. В нашем первом посте мы продемонстрировали процесс — и значительный потенциал — итеративного анализа и оптимизации модели PyTorch с использованием PyTorch Profiler и TensorBoard. В этом посте мы сосредоточимся на конкретном типе проблемы с производительностью, которая особенно распространена в PyTorch из-за использования нетерпеливого выполнения: зависимости от ЦП для частей выполнения модели. Выявление наличия и источника таких проблем может быть довольно сложным и часто требует использования специального анализатора производительности. В этом посте мы поделимся некоторыми советами по выявлению таких проблем с производительностью при использовании PyTorch Profiler и плагина PyTorch Profiler TensorBoard.

Плюсы и минусы нетерпеливого исполнения

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

Пример игрушки

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

Начнем с определения простой модели классификации. Его архитектура не имеет значения для этого поста.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim
import torch.profiler
import torch.utils.data
import torchvision.models
import torchvision.transforms as T
from torchvision.datasets.vision import VisionDataset
import numpy as np
from PIL import Image


# sample model
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 8, 3, padding=1)
        self.conv2 = nn.Conv2d(8, 12, 3, padding=1)
        self.conv3 = nn.Conv2d(12, 16, 3, padding=1)
        self.conv4 = nn.Conv2d(16, 20, 3, padding=1)
        self.conv5 = nn.Conv2d(20, 24, 3, padding=1)
        self.conv6 = nn.Conv2d(24, 28, 3, padding=1)
        self.conv7 = nn.Conv2d(28, 32, 3, padding=1)
        self.conv8 = nn.Conv2d(32, 10, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = self.pool(F.relu(self.conv4(x)))
        x = self.pool(F.relu(self.conv5(x)))
        x = self.pool(F.relu(self.conv6(x)))
        x = self.pool(F.relu(self.conv7(x)))
        x = self.pool(F.relu(self.conv8(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        return x

Далее мы определяем довольно стандартную функцию кросс-энтропийных потерь. Эта функция потерь будет в центре нашего обсуждения.

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def weighted_nll(pred, target, weight):
    assert target.max() < 10
    nll = -pred[range(target.shape[0]), target]
    nll = nll * weight[target]
    nll = nll / weight[target].sum()
    sum_nll = nll.sum()
    return sum_nll

# custom loss definition
class CrossEntropyLoss(nn.Module):
    def forward(self, input, target):
        pred = log_softmax(input)
        loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda())
        return loss

Наконец, мы определяем набор данных и цикл обучения:

# dataset with random images that mimics the properties of CIFAR10
class FakeCIFAR(VisionDataset):
    def __init__(self, transform):
        super().__init__(root=None, transform=transform)
        self.data = np.random.randint(low=0,high=256,size=(10000,32,32,3),dtype=np.uint8)
        self.targets = np.random.randint(low=0,high=10,size=(10000),dtype=np.uint8).tolist()

    def __getitem__(self, index):
        img, target = self.data[index], self.targets[index]
        img = Image.fromarray(img)
        if self.transform is not None:
            img = self.transform(img)
        return img, target

    def __len__(self) -> int:
        return len(self.data)

transform = T.Compose(
    [T.Resize(256),
     T.PILToTensor()])

train_set = FakeCIFAR(transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=1024,
                               shuffle=True, num_workers=8, pin_memory=True)

device = torch.device("cuda:0")
model = Net().cuda(device)
criterion = CrossEntropyLoss().cuda(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()


# training loop wrapped with profiler object
with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=4, active=3, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler(’./log/example’),
        record_shapes=True,
        profile_memory=True,
        with_stack=True
) as prof:
    for step, data in enumerate(train_loader):
        inputs = data[0].to(device=device, non_blocking=True)
        labels = data[1].to(device=device, non_blocking=True)
        inputs = (inputs.to(torch.float32) / 255. - 0.5) / 0.5
        if step >= (1 + 4 + 3) * 1:
            break
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        prof.step()

Опытный разработчик PyTorch, возможно, уже заметил, что наш пример содержит ряд неэффективных строк кода в функции потерь. В то же время в этом нет ничего явно плохого, и такие типы неэффективности не являются чем-то необычным. Если вы хотите проверить свои навыки PyTorch, посмотрите, сможете ли вы найти три проблемы с нашей реализацией потери перекрестной энтропии, прежде чем читать дальше. В следующих разделах мы предположим, что не смогли найти эти проблемы самостоятельно, и покажем, как мы можем использовать PyTorch Profiler и связанный с ним плагин TensorBoard для их выявления.

Как и в нашем предыдущем посте, мы итеративно проведем эксперимент, выявим проблемы с производительностью и попытаемся их исправить. Мы проведем наши эксперименты на экземпляре Amazon EC2 g5.2xlarge (содержащем графический процессор NVIDIA A10G и 8 виртуальных ЦП) и используя официальный образ AWS PyTorch 2.0 Docker. Наш выбор тренировочной среды был несколько произвольным, и его не следует рассматривать как одобрение какого-либо из ее компонентов.

Начальные результаты производительности

На изображении ниже мы показываем вкладку «Обзор» отчета о производительности скрипта выше.

Как мы видим, использование нашего графического процессора находится на относительно высоком уровне 92,04%, а время шага составляет 216 миллисекунд. (Как и в нашем предыдущем посте, Обзор в torch-tb-profiler версии 0.4.1 суммирует время всех трех шагов обучения.) Только из этого отчета вы можете не подумать, что с нашей моделью что-то не так. Однако представление Trace View отчета о производительности рассказывает совершенно другую историю:

Как было отмечено выше, прямой проход нашей кросс-энтропийной потери одной занимает 211 из 216 миллисекунд шага обучения! Это явный признак того, что что-то не так. Наша функция потерь содержит небольшое количество вычислений по сравнению с моделью и, безусловно, не должна учитывать 98% времени шага. Присмотревшись к стеку вызовов, мы можем увидеть несколько вызовов функций, которые усиливают наши подозрения, в том числе «to», «copy_» и «cudaStreamSynchronize». Эта комбинация обычно указывает на то, что данные копируются из ЦП в ГП, а это не то, что мы хотим, чтобы происходило в середине нашего расчета потерь. В этом случае наша проблема с производительностью также связана с кратковременным снижением использования графического процессора, как показано на изображении. Однако это не всегда так. Часто провалы в использовании GPU не связаны с проблемой производительности или вообще не видны.

Теперь мы знаем, что у нас есть проблема с производительностью в нашей функции потерь и что она, вероятно, связана с копированием тензоров с хоста на графический процессор. Однако этого может быть недостаточно, чтобы определить точную строку кода, вызывающую проблему. Чтобы облегчить наш поиск, мы обернем каждую строку кода менеджером контекста с пометкой torch.profiler.record_function и повторно запустим анализ профилирования.

# custom loss definition
class CrossEntropyLoss(nn.Module):
    def forward(self, input, target):
        with torch.profiler.record_function('log_softmax'):
            pred = log_softmax(input)
        with torch.profiler.record_function('define_weights'):
            weights = torch.Tensor([0.1]*10).cuda()
        with torch.profiler.record_function('weighted_nll'):
            loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda())
        return loss

Добавление меток помогает нам идентифицировать определение веса или, точнее, копирование весов в графический процессор как проблемную строку кода.

Оптимизация № 1. Удалите избыточные копии хост-процессор на этапе обучения.

Как только мы определили нашу первую проблему, ее исправление довольно тривиально. В приведенном ниже блоке кода мы копируем наш весовой вектор в GPU один раз в функции loss init:

class CrossEntropyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.Tensor([0.1]*10).cuda()

    def forward(self, input, target):
        with torch.profiler.record_function('log_softmax'):
            pred = log_softmax(input)
        with torch.profiler.record_function('weighted_nll'):
            loss = weighted_nll(pred, target, self.weight)
        return loss

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

К сожалению, наша первая оптимизация оказала очень незначительное влияние на время шага. Если мы посмотрим на отчет Trace View, мы увидим, что у нас есть новая серьезная проблема с производительностью, которую нам нужно решить.

Наш новый отчет указывает на проблему, связанную с нашей функцией weighted_nll. Как и прежде, мы использовали torch.profiler.record_function для определения проблемной строки кода. В данном случае это вызов assert.

def weighted_nll(pred, target, weight):
    with torch.profiler.record_function('assert'):
        assert target.max() < 10
    with torch.profiler.record_function('range'):
        r = range(target.shape[0])
    with torch.profiler.record_function('index'):
        nll = -pred[r, target]
    with torch.profiler.record_function('nll_calc'):
        nll = nll * weight[target]
        nll = nll/ weight[target].sum()
        sum_nll = nll.sum()
    return sum_nll

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

Более тщательный анализ стека вызовов показывает вызовы «item», «_local_scalar_dense» и «cudaMemcpyAsync». Часто это указывает на то, что данные копируются с графического процессора на хост. Действительно, наш вызов assert, который выполняется на ЦП, требует доступа к целевомутензору, находящемуся на графическом процессоре, что приводит к крайне неэффективному копированию данных.

Оптимизация № 2. Удалите избыточные копии GPU-to-host на этапе обучения.

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

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

Анализ представления трассировки знакомит нас со следующей проблемой производительности:

И снова мы видим, что наша предыдущая оптимизация выявила новую серьезную проблему с производительностью, на этот раз при индексации нашего тензора pred. Индексы определяются тензорами r и target. В то время как тензор target уже находится в графическом процессоре, тензор r, определенный в предыдущей строке, отсутствует. Это снова приводит к неэффективному копированию данных с хоста на GPU.

Оптимизация №3: Заменить диапазон на torch.arange

Функция Python range выводит список на ЦП. Наличие любого списка на вашем этапе обучения должно быть красным флажком. В приведенном ниже блоке кода мы заменяем использование range на torch.arange и настраиваем его для создания выходного тензора непосредственно на GPU:

def weighted_nll(pred, target, weight):
    with torch.profiler.record_function('range'):
        r = torch.arange(target.shape[0], device="cuda:0")
    with torch.profiler.record_function('index'):
        nll = -pred[r, target]
    with torch.profiler.record_function('nll_calc'):
        nll = nll * weight[target]
        nll = nll/ weight[target].sum()
        sum_nll = nll.sum()
    return sum_nll

Результаты этой оптимизации показаны ниже:

Сейчас мы говорим!! Наше время шага сократилось до 5,8 миллисекунд, что означает увеличение производительности на колоссальные 3700%.

Обновленный вид трассировки показывает, что функция потерь снизилась до очень разумных 0,5 миллисекунд.

Но есть еще возможности для улучшения. Давайте подробнее рассмотрим представление трассировки функции weighted_nll, которая выполняет большую часть расчета потерь.

Из трассировки видно, что функция формируется из нескольких небольших блоков, каждый из которых в конечном итоге сопоставляется с отдельным ядром CUDA, которое загружается в GPU с помощью вызова CudaLaunchKernel. В идеале мы хотели бы уменьшить общее количество ядер графического процессора, чтобы уменьшить количество взаимодействий между процессором и графическим процессором. Один из способов сделать это — по возможности предпочесть операторы PyTorch более высокого уровня, такие как torch.nn.NLLLoss. Предполагается, что такие функции объединяют базовые операции, что требует меньшего количества общих ядер.

Оптимизация №4: заменить пользовательский NLL на torch.nn.NLLLoss

Блок кода ниже содержит наше обновленное определение потерь, в котором теперь используется torch.nn.NLLLoss.

class CrossEntropyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.Tensor([0.1]*10).cuda()

    def forward(self, input, target):
        pred = log_softmax(input)
        nll = torch.nn.NLLLoss(self.weight)
        loss = nll(pred, target)
        return loss

Здесь мы позволили себе ввести еще одну распространенную ошибку, которую мы и продемонстрируем.

Использование функции более высокого уровня еще больше сокращает время шага до 5,3 миллисекунды (по сравнению с 5,8).

Однако, если мы внимательно посмотрим на Trace View, то увидим, что значительная часть функции потерь теперь тратится на инициализацию объекта torch.nn.NLLLoss!

Оглядываясь назад на нашу функцию потерь, мы видим, что мы инициализируем новый объект NLLLoss на каждой итерации шага обучения. Естественно, инициализация объекта происходит на ЦП, и хотя (в нашем случае) это относительно быстро, мы хотели бы избежать этого на этапе обучения.

Оптимизация № 5: воздержитесь от инициализации объектов на этапе обучения

В приведенном ниже блоке кода мы изменили нашу реализацию потерь, чтобы в функции init был создан единственный экземпляр torch.nn.NLLLoss.

class CrossEntropyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.Tensor([0.1]*10).cuda()
        self.nll = torch.nn.NLLLoss(self.weight) 

    def forward(self, input, target):
        pred = log_softmax(input)
        loss = self.nll(pred, target)
        return loss

Результаты показывают дальнейшее улучшение времени шага, которое теперь составляет 5,2 миллисекунды.

Оптимизация №6: Используйте torch.nn.CrossEntropyLoss вместо пользовательского лосса

PyTorch включает встроенный torch.nn.CrossEntropyLoss, который мы сейчас оцениваем и сравниваем с нашей собственной реализацией потерь.

criterion = torch.nn.CrossEntropyLoss().cuda(device)

Результирующее время шага стало новым минимумом в 5 миллисекунд для общего прироста производительности на 4200% (по сравнению с 216 миллисекундами, с которых мы начали).

Улучшение производительности прямого прохода расчета потерь еще более заметно: с начальной точки в 211 миллисекунд мы снизились до 79 микросекундсекунд( !!), как показано ниже:

Оптимизация № 7: компиляция функции потерь

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

criterion = torch.compile(torch.nn.CrossEntropyLoss().cuda(device))

На изображении ниже показан результат просмотра трассировки этого эксперимента.

Первое, что мы видим, это появление терминов, содержащих OptimizedModule и dynamo, которые указывают на использование torch.compile. Мы также можем видеть, что на практике компиляция модели не уменьшила количество ядер, загруженных функцией потерь, что означает, что она не выявила никаких возможностей для дополнительного слияния ядер. Фактически, в нашем случае компиляция потерь фактически привела к увеличению времени прямого прохода функции потерь с 79 до 154 микросекунд. Похоже, что CrossEntropyLoss недостаточно содержателен, чтобы извлечь выгоду из этой оптимизации.

Вам может быть интересно, почему мы не можем просто применить компиляцию torch к нашей начальной функции потерь и полагаться на нее для оптимальной компиляции нашего кода. Это может избавить вас от всех хлопот пошаговой оптимизации, которые мы описали выше. Проблема с этим подходом заключается в том, что хотя компиляция PyTorch 2.0 (на момент написания этой статьи) действительно оптимизирует определенные типы кроссоверов GPU-CPU, некоторые типы приведут к сбою компиляции графа, а другие приведут к созданию нескольких небольшие графики, а не один большой. Последняя категория вызывает разрывы графика, что существенно ограничивает возможности функции torch.compile по повышению производительности. (Один из способов решить эту проблему — вызвать torch.compile с флагом fullgraph, установленным в True.) См. наш предыдущий пост для получения более подробной информации об использовании этой опции.

Полученные результаты

В таблице ниже мы суммируем результаты проведенных нами экспериментов:

Наши последовательные оптимизации привели к умопомрачительному увеличению производительности на 4143%! Напомним, что мы начали с довольно невинно выглядящей функции потерь. Без глубокого анализа поведения нашего приложения мы, возможно, никогда не узнали бы, что что-то не так, и продолжали бы жить, заплатив в 41 раз (!!) больше, чем нам нужно.

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

Выводы

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

Советы по кодированию

То, как вы реализуете прямой проход вашей модели, может оказать существенное влияние на ее производительность. Здесь мы перечислим лишь несколько рекомендаций, основанных на примере, который мы рассмотрели в этом посте.

Избегайте инициализации постоянных тензоров в прямом проходе. Вместо этого сделайте это в конструкторе.
Избегайте использования asserts для тензоров, находящихся на графическом процессоре, в прямом проходе. Либо переместите их в конвейер ввода данных, либо проверьте, есть ли в PyTorch какие-либо встроенные методы для выполнения необходимой вам проверки данных.
Избегайте использования списков. Проверьте, может ли использование torch.arange для создания тензора непосредственно на устройстве быть лучшей альтернативой.
Используйте операторы PyTorch, такие как torch.nn.NLLLoss и torch.nn.CrossEntropyLoss, вместо создания собственных реализаций потерь.
Избегайте инициализации объектов на прямом проходе. Вместо этого сделайте это в конструкторе.
При необходимости рассмотрите возможность использования torch.compile.

Советы по анализу производительности

Как мы продемонстрировали, Trace View плагина Tensorboard PyTorch Profiler имел решающее значение для выявления проблем с производительностью в нашей модели. Ниже мы резюмируем некоторые основные выводы из нашего примера:

Высокая загрузка графического процессора НЕ обязательно является признаком того, что ваш код работает оптимально.
Обратите внимание на части кода, которые занимают больше времени, чем ожидалось.
Используйте torch.profiler.record_function, чтобы выявить проблемы с производительностью.
Провалы в использовании графического процессора не обязательно связаны с источником проблемы с производительностью.
Следите за непреднамеренными копиями данных с хоста на GPU. Обычно они идентифицируются вызовами «to», «copy_» и «cudaStreamSynchronize», которые можно найти в представлении трассировки.
Следите за непреднамеренными копиями данных с графического процессора на хост. Обычно они идентифицируются вызовами «item» и «cudaStreamSynchronize», которые можно найти в представлении трассировки.

Краткое содержание

В этом посте мы сосредоточились на проблемах производительности в обучающих приложениях, возникающих из-за избыточного взаимодействия между ЦП и ГП во время прямого прохода шага обучения. Мы продемонстрировали, как можно использовать анализаторы производительности, такие как PyTorch Profiler и связанный с ним подключаемый модуль TensorBoard, для выявления таких проблем и значительного повышения производительности.

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

Анализ и оптимизация производительности модели PyTorch — часть 2