Решение узких мест на входном конвейере данных с помощью PyTorch Profiler и TensorBoard

'Solving bottlenecks on the input data pipeline with PyTorch Profiler and TensorBoard

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

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

1. Выбор образца обучения с соотношением вычислений между CPU и GPU, наиболее подходящим для вашей рабочей нагрузки (например, см. нашу предыдущую публикацию с советами по выбору наилучшего типа экземпляра для вашей задачи машинного обучения),
2. Улучшение баланса нагрузки между CPU и GPU путем перемещения части предварительной обработки CPU на GPU,
3. Перенос некоторых вычислений CPU на вспомогательные устройства с CPU-воркерами (например, здесь).

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

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

Большое спасибо Йицхаку Леви за его вклад в эту публикацию.

• Предсказание зарплат NBA с помощью машинного обучения
• 5 причин, чтобы рассмотреть участие в интенсивной программе по науке о данных перед поступлением в высшее учебное заведение
• Архитектура данных и Теорема CAP где происходит столкновение?

Игрушечная модель

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

from math import log2import torchimport torch.nn as nnimport torch.nn.functional as Fimg_size = 256num_classes = 10hidden_size = 30# игрушечная модель классификации CNNclass Net(nn.Module):    def __init__(self, img_size=img_size, num_classes=num_classes):        super().__init__()        self.conv_in = nn.Conv2d(3, hidden_size, 3, padding='same')        num_hidden = int(log2(img_size))        hidden = []        for i in range(num_hidden):            hidden.append(nn.Conv2d(hidden_size, hidden_size, 3, padding='same'))            hidden.append(nn.ReLU())            hidden.append(nn.MaxPool2d(2))        self.hidden = nn.Sequential(*hidden)        self.conv_out = nn.Conv2d(hidden_size, num_classes, 3, padding='same')    def forward(self, x):        x = F.relu(self.conv_in(x))        x = self.hidden(x)        x = self.conv_out(x)        x = torch.flatten(x, 1)        return x

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

import numpy as npfrom PIL import Imagefrom torchvision.datasets.vision import VisionDatasetinput_img_size = [533, 800]class FakeDataset(VisionDataset):    def __init__(self, transform):        super().__init__(root=None, transform=transform)        size = 10000        self.img_files = [f'0.jpg' for i in range(size)]        self.targets = np.random.randint(low=0,high=num_classes,                                         size=(size),dtype=np.uint8).tolist()    def __getitem__(self, index):        img_file, target = self.img_files[index], self.targets[index]        with torch.profiler.record_function('PIL open'):            img = Image.open(img_file)        if self.transform is not None:            img = self.transform(img)        return img, target    def __len__(self):        return len(self.img_files)

Обратите внимание, что мы обернули считыватель файла в менеджер контекста torch.profiler.record_function.

Наша конвейер ввода данных включает следующие преобразования изображения:

1. PILToTensor преобразует изображение PIL в Tensor PyTorch.
2. RandomCrop возвращает обрезку 256×256 с случайным смещением в изображении.
3. RandomMask – это пользовательское преобразование, которое создает случайную булеву маску размером 256×256 и применяет ее к изображению. Преобразование включает операцию расширения на четырех соседей маски.
4. ConvertColor – это пользовательское преобразование, которое преобразует формат изображения из RGB в YUV.
5. Scale – это пользовательское преобразование, которое масштабирует пиксели в диапазон [0,1].

class RandomMask(torch.nn.Module):    def __init__(self, ratio=0.25):        super().__init__()        self.ratio=ratio    def dilate_mask(self, mask):        # выполняем расширение маски на четырех соседей        with torch.profiler.record_function('dilation'):            from scipy.signal import convolve2d            dilated = convolve2d(mask, [[0, 1, 0],                                     [1, 1, 1],                                     [0, 1, 0]], mode='same').astype(bool)        return dilated    def forward(self, img):        with torch.profiler.record_function('random'):            mask = np.random.uniform(size=(img_size, img_size)) < self.ratio        dilated_mask = torch.unsqueeze(torch.tensor(self.dilate_mask(mask)),0)        dilated_mask = dilated_mask.expand(3,-1,-1)        img[dilated_mask] = 0.        return img    def __repr__(self):        return f"{self.__class__.__name__}(ratio={self.ratio})"class ConvertColor(torch.nn.Module):    def __init__(self):        super().__init__()        self.A=torch.tensor(            [[0.299, 0.587, 0.114],             [-0.16874, -0.33126, 0.5],             [0.5, -0.41869, -0.08131]]        )        self.b=torch.tensor([0.,128.,128.])    def forward(self, img):        img = img.to(dtype=torch.get_default_dtype())        img = torch.matmul(self.A,img.view([3,-1])).view(img.shape)        img = img + self.b[:,None,None]        return img    def __repr__(self):        return f"{self.__class__.__name__}()"class Scale(object):    def __call__(self, img):        return img.to(dtype=torch.get_default_dtype()).div(255)    def __repr__(self):        return f"{self.__class__.__name__}()"

Мы объединяем преобразования с помощью класса Compose, который мы немного модифицировали, чтобы включить менеджер контекста torch.profiler.record_function вокруг каждого из вызовов преобразований.

import torchvision.transforms as Tclass CustomCompose(T.Compose):    def __call__(self, img):        for t in self.transforms:            with torch.profiler.record_function(t.__class__.__name__):                img = t(img)        return imgtransform = CustomCompose(    [T.PILToTensor(),     T.RandomCrop(img_size),     RandomMask(),     ConvertColor(),     Scale()])

В блоке кода ниже мы определяем набор данных и загрузчик данных. Мы настраиваем DataLoader для использования пользовательской функции collate, в которой мы оборачиваем функцию collate по умолчанию в менеджер контекста torch.profiler.record_function.

train_set = FakeDataset(transform=transform)def custom_collate(batch):    from torch.utils.data._utils.collate import default_collate    with torch.profiler.record_function('collate'):        batch = default_collate(batch)    image, label = batch    return image, labeltrain_loader = torch.utils.data.DataLoader(train_set, batch_size=256,                                           collate_fn=custom_collate,                                           num_workers=4, pin_memory=True)

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

from statistics import mean, variancefrom time import timedevice = torch.device("cuda:0")model = Net().cuda(device)criterion = nn.CrossEntropyLoss().cuda(device)optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)model.train()t0 = time()times = []with torch.profiler.profile(    schedule=torch.profiler.schedule(wait=10, warmup=2, active=10, repeat=1),    on_trace_ready=torch.profiler.tensorboard_trace_handler('/tmp/prof'),    record_shapes=True,    profile_memory=True,    with_stack=True) as prof:    for step, data in enumerate(train_loader):        with torch.profiler.record_function('h2d copy'):            inputs, labels = data[0].to(device=device, non_blocking=True), \                             data[1].to(device=device, non_blocking=True)        if step >= 40:            break        outputs = model(inputs)        loss = criterion(outputs, labels)        optimizer.zero_grad(set_to_none=True)        loss.backward()        optimizer.step()        prof.step()        times.append(time()-t0)        t0 = time()print(f'average time: {mean(times[1:])}, variance: {variance(times[1:])}')

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

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

Среднее время выполнения шага, указанное в нашем скрипте, составляет 1,3 секунды, а средняя загрузка GPU очень низкая – 18,21%. На изображении ниже мы запечатлели результаты производительности, отображаемые в Trace View плагина TensorBoard:

Мы видим, что каждый четвертый шаг обучения включает длительный (~5,5 секунд) период загрузки данных, во время которого GPU полностью простаивает. Причина этого заключается в количестве рабочих процессов DataLoader, которое мы выбрали – четыре. Каждый четвертый шаг мы видим, что все рабочие процессы заняты созданием образцов для следующей пакетной выборки, в то время как GPU ждет. Это явный признак узкого места в конвейере ввода данных. Вопрос в том, как мы его анализируем? Затрудняет дело тот факт, что множество маркеров record_function, которые мы вставили в код, нигде не отображаются в трассировке профиля.

Использование нескольких рабочих процессов в DataLoader является критически важным для оптимизации производительности. К сожалению, это также усложняет процесс профилирования. Хотя существуют профилировщики, поддерживающие анализ многопроцессорной работы (например, посмотрите VizTracer), подход, который мы выберем в этом сообщении, заключается в запуске, анализе и оптимизации нашей модели в режиме однопроцессорного режима (т.е. без рабочих процессов DataLoader), а затем применении оптимизаций к многопроцессорному режиму. Признаться, оптимизация скорости отдельной функции не гарантирует, что множественные (параллельные) вызовы этой же функции также будут иметь выгоды. Однако, как мы увидим в этом сообщении, такая стратегия позволит нам выявить и решить некоторые основные проблемы, которые мы не смогли идентифицировать иначе, и, по крайней мере, в отношении здесь обсуждаемых вопросов мы обнаружим сильную корреляцию между влиянием на производительность в двух режимах. Но прежде чем мы применим эту стратегию, давайте настроим выбор количества рабочих процессов.

Оптимизация 1: Настройка стратегии многопроцессорности

Определение оптимального числа потоков или процессов в многопроцессорном/многопоточном приложении, таком как наше, может быть сложной задачей. С одной стороны, если мы выберем слишком низкое число, мы можем неэффективно использовать ресурсы ЦП. С другой стороны, если мы выберем слишком высокое число, мы рискуем затратами, нежелательной ситуацией, при которой операционная система тратит большую часть времени на управление многопоточностью/многопроцессорностью, а не на выполнение нашего кода. В случае рабочей нагрузки обучения PyTorch рекомендуется протестировать разные варианты настройки num_workers для DataLoader. Хорошим местом для начала является установка числа на основе числа ЦП на хосте (например, num_workers:=num_cpus/num_gpus). В нашем случае на Amazon EC2 g5.2xlarge есть восемь виртуальных процессоров, и, фактически, увеличение числа рабочих процессов DataLoader до восьми приводит к незначительному улучшению среднего времени выполнения шага до 1,17 секунды (увеличение на 11%).

Важно обратить внимание на другие, менее очевидные настройки конфигурации, которые могут влиять на количество потоков или процессов, используемых конвейером ввода данных. Например, opencv-python, библиотека, часто используемая для предварительной обработки изображений в задачах компьютерного зрения, включает функцию cv2.setNumThreads(int) для установки количества потоков.

На изображении ниже мы запечатлели часть Trace View при запуске скрипта с параметром num_workers, установленным в ноль.

Запуск скрипта таким образом позволяет нам увидеть метки record_function, которые мы установили, и идентифицировать операцию RandomMask transform, или, более конкретно, нашу функцию dilation, как самую затратную операцию при извлечении каждого отдельного образца.

Оптимизация 2: Оптимизация функции дилатации

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

В приведенном ниже блоке кода мы предлагаем альтернативную, более удобную для ЦП, реализацию функции дилатации, использующую булевые операции вместо свертки:

    def расширить_маску(self, маска):        # выполнять дилатацию на маске с 4 соседями        with torch.profiler.record_function('dilation'):            padded = np.pad(mask, [(1,1),(1,1)])            dilated = padded[0:-2,1:-1] | padded[1:-1,1:-1] | padded[2:,1:-1] | padded[1:-1,0:-2]| padded[1:-1,2:]        return расширенный

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

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

При более детальном рассмотрении функции PILToTensor (см. здесь) обнаруживаются три подлежащие операции:

1. Загрузка изображения PIL – из-за свойства ленивой загрузки Image.open, изображение загружается здесь.
2. Изображение PIL преобразуется в массив numpy.
3. Массив numpy преобразуется в тензор PyTorch.

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

Оптимизация 3: Изменение порядка преобразований

К счастью, преобразование RandomCrop можно применить непосредственно к изображению PIL, что позволяет нам применить уменьшение размера изображения в качестве первой операции в нашей конвейерной линии:

transform = CustomCompose(    [T.RandomCrop(img_size),     T.PILToTensor(),     RandomMask(),     ConvertColor(),     Scale()])

После этой оптимизации наше время шага снижается до 0,72 секунды, что является дополнительной оптимизацией на 8%. Ниже приведена захваченная трассировка, показывающая, что преобразование RandomCrop теперь является доминирующей операцией:

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

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

Оптимизация 4: Применение пакетных преобразований

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

def batch_transform(img):    img = img.to(dtype=torch.get_default_dtype())    A = torch.tensor(        [[0.299, 0.587, 0.114],         [-0.16874, -0.33126, 0.5],         [0.5, -0.41869, -0.08131]]    )    b = torch.tensor([0., 128., 128.])    A = torch.broadcast_to(A, ([img.shape[0],3,3]))    t_img = torch.bmm(A,img.view(img.shape[0],3,-1))    t_img = t_img + b[None,:, None]    return t_img.view(img.shape)/255def custom_collate(batch):    from torch.utils.data._utils.collate import default_collate    with torch.profiler.record_function('collate'):        batch = default_collate(batch)    image, label = batch    with torch.profiler.record_function('batch_transform'):        image = batch_transform(image)    return image, label

Результатом этого изменения является небольшое увеличение времени шага до 0,75 секунды. Хотя это бесполезно в случае нашей игрушечной модели, возможность применения определенных операций в виде пакетных преобразований, а не преобразований на каждом примере, может оптимизировать определенные рабочие нагрузки.

Результаты

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

Резюме

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

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