Практика: мой опыт интеграции более 50 нейронных сетей в один проект

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

Проект ориентирован на создание и редактирование видео, изображений и аудио с применением нейронных сетей. Часто разные методы могут выполнять схожие задачи. Так как я интегрировал решения с открытым исходным кодом, оптимизировал их и добавлял новый функционал, ключевой задачей стало обеспечение единства методов. Например, такие функции, как замена лица, синхронизация губ и анимация портретов, требуют распознавания лица. В моем проекте за эту задачу отвечает одна модель, а не несколько разных методов, как в исходных решениях. Поэтому все 50+ моделей распределены так, что каждая отвечает за своё уникальное направление, без дублирования.

Одна модель = Одна задача

В процессе разработки я принципиально отказался от TensorFlow и связанных с ним решений, сосредоточившись исключительно на PyTorch и ONNX Runtime.

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

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

Итак, начнем

Лайфхак 1

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

Лайфхак 2

Очередь. Мое приложение основано на Flask, поэтому пользователь не ожидает окончания обработки и может запускать сколько угодно задач, тем самым загружая память. В результате я искусственно создаю задержку между задачами с случайным значением, чтобы избежать одновременного запуска двух и более задач. Это связано с Лайфхаком 3.

Лайфхак 3

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

import torch
import psutil

def get_vram_gb(device="cuda"):
    if torch.cuda.is_available():
        properties = torch.cuda.get_device_properties(device)  # Get the values ​​for a specific GPU, which is our device
        total_vram_gb = properties.total_memory / (1024 ** 3)
        available_vram_gb = (properties.total_memory - torch.cuda.memory_allocated()) / (1024 ** 3)
        busy_vram_gb = total_vram_gb - available_vram_gb
        return total_vram_gb, available_vram_gb, busy_vram_gb
    return 0, 0, 0


def get_ram_gb():
    mem = psutil.virtual_memory()
    total_ram_gb = mem.total / (1024 ** 3)
    available_ram_gb = mem.available / (1024 ** 3)
    busy_ram_gb = total_ram_gb - available_ram_gb
    return total_ram_gb, available_ram_gb, busy_ram_gb

Лайфхак 4

Вместе с отложенным запуском я использую проверки на самую распространенную ошибку: “CUDA out of memory”. Идея заключается в том, что если мы получаем сообщение о нехватке памяти, нам нужно очистить память от ненужных данных и запустить процесс заново.

min_delay = 20
max_delay = 180
try:
    # Launch the method with a neural network
except RuntimeError as err:
    if 'CUDA out of memory' in str(err):
        # Clear memory
        sleep(random.randint(min_delay, max_delay))
        # Clear memory again
        # Launch the method again
    else:
        raise err

К этой части мы ещё вернемся, поскольку недостаточно просто выполнить `# Clear cache`, всё должно быть немного иначе.

Лайфхак 5

Backend моей программы состоит из модулей, которые классифицируются по следующим признакам: изменение видео или изображения, генерация видео и изображений, изменение аудио — т.е. по свойству модели. И также по признаку: модель обрабатывает задачи для frontend или backend, т.е. результат работы модели необходимо вернуть мгновенно пользователю (сегментация, txt2img и img2img) или как выполненную крупную задачу. Мы не говорим про модели, которые работают на frontend, используя:

await ort.InferenceSession.create(MODEL_DIR).then(console.log("Model loaded"))

Следовательно, мне необходимо загружать модели для быстрого возврата ответа в память и держать их там, не позволяя разным пользователям одновременно использовать одну модель (Лайфхак 1) и не использовать их для задач с долгой обработкой, чтобы не нарушить Лайфхак 1.

Лайфхак 6

Модели для длительной обработки иногда бывают очень требовательными, и в зависимости от видеопамяти, такая модель может полностью её использовать. В плане оптимизации очень невыгодно каждый раз загружать и выгружать такие модели, хотя иногда это, к сожалению, приходится делать. Часто с такими моделями используются микро модели, которые занимают в памяти немного места, но их загрузка и выгрузка требует времени. При запуске задач мы группируем их по методам длительной обработки, и задачи из одной группы обрабатываются на маленьких моделях, создавая очередь перед загрузкой в одну большую модель. Помните Лайфхаки 3 и 4? У нас есть два метода: измерить, сколько такая модель потребляет памяти, или запустить её, чтобы получить ошибку “CUDA out of memory” и очистить кэш.

Очистка от кэша

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

if torch.cuda.is_available():  # If CUDA is available, because the application can work without CUDA
	torch.cuda.empty_cache() # Frees unused memory in the CUDA cache
	torch.cuda.ipc_collect() # Performs garbage collection on CUDA objects accessed via IPC (interprocess communication)
gc.collect() # Calls Python's garbage collector to free memory occupied by unused objects

Лайфхак 7

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

del ...

Лайфхак 8

Модели можно загружать по слоям на GPU и CPU, либо на несколько GPU, но при этом элементы одного слоя должны находиться на одном GPU. Такой подход применяется при малом количестве видеопамяти и используется в генерации изображений и видео, но не ограничивается этим.

device_map = {
    'encoder.layer.0': 'cuda:0',
    'encoder.layer.1': 'cuda:1',
    'decoder.layer.0': 'cuda:0',
    'decoder.layer.1': 'cuda:1',
}
# Or
device_map = {
    'encoder.layer.0': 'cuda',
    'encoder.layer.1': 'cpu',
    'decoder.layer.0': 'cuda',
    'decoder.layer.1': 'cpu',
}

Лайфхак 9

Не забывайте использовать enable_xformers_memory_efficient_attention(), если пайплайн модели это поддерживает. В документациях описаны и другие методы, такие как enable_model_cpu_offload(), enable_vae_tiling(), enable_attention_slicing(). У меня они работают при рестайлинге видео, а для генерации изображений используются совсем другие методы:

if vram < 12:
    pipe.enable_sequential_cpu_offload()
    print("VRAM below 12 GB: Using sequential CPU offloading for memory efficiency. Expect slower generation.")
elif vram < 20:
    print("VRAM between 12-20 GB: Medium generation speed enabled.")
elif vram < 30:
    # Load essential modules to GPU
    for module in [pipe.vae, pipe.dit, pipe.text_encoder]:
        module.to("cuda")
    cpu_offloading = False
    print("VRAM between 20-30 GB: Sufficient memory for faster generation.")
else:
    # Maximize performance by disabling memory-saving options
    for module in [pipe.vae, pipe.dit, pipe.text_encoder]:
        module.to("cuda")
    cpu_offloading = False
    save_memory = False
    print("VRAM above 30 GB: Maximum speed enabled for generation.")

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

Лайфхак 10

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

Лайфхак 11

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

Лайфхак 12

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

Лайфхак 13

Давайте поговорим о совместимости версий библиотек, особенно таких, как torch, torchvision, torchaudio и xformers. Важно, чтобы они были совместимы между собой и с вашей версией CUDA. Как мы поступаем?

Первое — проверяем версию своего CUDA:

nvcc -V

Второе — заходим на сайт PyTorch, чтобы ознакомиться с совместимостью версий: PyTorch Previous Versions или на страницу загрузки, где cu118 — это ваша версия CUDA. Обратите внимание, что ваша версия CUDA может работать с более старыми версиями torch. Например, CUDA 12.6 может работать с torch версии, совместимой с cu118.

Я заметил, что torch и torchaudio часто имеют одинаковые версии, например, 2.4.1, в то время как версия torchvision может отличаться, как, например, 0.19.1. Таким образом, можно определить, что torch и torchaudio версии 2.2.2 работают с torchvision 0.17.2. Чувствуете зависимость?

Дополнительно вы можете загружать файлы .whl по ссылке и даже распаковывать их самостоятельно. Для меня соблюдение версий критически важно, так как программа устанавливается через установщик, и для пользователей Windows, при первом включении загружаются torch, torchaudio и torchvision в зависимости от их выбора, с индикацией статуса загрузки, а потом распаковывает.

Третье — необходимо убедиться, что xformers также совместим. Для этого посетите репозиторий xformers на GitHub и внимательно ознакомьтесь с тем, с какой версией torch и CUDA будет работать xformers, так как поддержка старых версий может быть отменена, в том числе для torch. Например, при использовании CUDA 11.8 вы ощутите пользу от xformers, особенно если ваше устройство имеет ограниченное количество видеопамяти.

Четвертое — это не обязательный шаг, но есть такая вещь, как flash-attn. Если вы решите её установить, вы можете сделать это быстрее, используя команду:

MAX_JOBS=4 pip install flash-attn

Где вы можете выбрать количество jobs, которое вам подходит. Я использую её следующим образом:

try:
    from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
    from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis
    from flash_attn.flash_attn_interface import flash_attn_varlen_func
except ImportError:
    flash_attn_func = None
    flash_attn_qkvpacked_func = None
    flash_attn_varlen_func = None

Лайфхак 14

Чтобы убедиться, что CUDA доступна в провайдерах ONNX Runtime, выполните следующий код:

access_providers = onnxruntime.get_available_providers()
if "CUDAExecutionProvider" in access_providers:
    provider = ["CUDAExecutionProvider"] if torch.cuda.is_available() and self.device == "cuda" else ["CPUExecutionProvider"]
else:
    provider = ["CPUExecutionProvider"]

Для новых версий CUDA 12.x, в отличие от более старой версии 11.8, вам также потребуется установить cuDNN 9.x на Linux (на Windows это может быть не обязательно). Обратите внимание, что иногда onnxruntime-gpu устанавливается без поддержки CUDA. Поэтому, когда мы убедимся, что версия torch совместима с CUDA, рекомендуется переустановить onnxruntime-gpu:

pip install -U onnxruntime-gpu

Лайфхак 15

Что делать, если некоторые модели работают только со старыми библиотеками, а другие — только с новыми? Я столкнулся с такой проблемой в gfpganer, где он требует старую версию torchvision, в то время как для генерации видео необходимы новые версии torch. В этом случае вы можете воспользоваться следующим подходом:

try:
    # Check if `torchvision.transforms.functional_tensor` and `rgb_to_grayscale` are missing
    from torchvision.transforms.functional_tensor import rgb_to_grayscale
except ImportError:
    # Import `rgb_to_grayscale` from `functional` if it’s missing in `functional_tensor`
    from torchvision.transforms.functional import rgb_to_grayscale

    # Create a module for `torchvision.transforms.functional_tensor`
    functional_tensor = types.ModuleType("torchvision.transforms.functional_tensor")
    functional_tensor.rgb_to_grayscale = rgb_to_grayscale

    # Add this module to `sys.modules` so other imports can access it
    sys.modules["torchvision.transforms.functional_tensor"] = functional_tensor

Таким образом, вы импортируете измененные методы для тех, которые исчезли в новых версиях. Это позволяет обеспечить совместимость между различными библиотеками и моделями.

Лайфхак 16

Обращайте внимание на предупреждения (Warning). Всегда следите за сообщениями типа Warning, в которых говорится о предстоящих изменениях в новых версиях библиотек. Ищите соответствующие строки кода в вашем проекте и добавляйте или изменяйте необходимые параметры. Это поможет избежать накопления несоответствий при обновлении до новых версий.

Когда увидел Warning в консоли

Лайфхак 17

Управление GPU в кластере. Если вы используете кластер из нескольких машин, помните, что вы не можете суммировать видеопамять от разных GPU. Однако, если видеокарты находятся в локальной сети, вы можете использовать управление GPU из одного контроллера. Для этого существуют библиотеки, такие как Ray. Обратите внимание, что суммирование видеопамяти не работает, за исключением случаев, когда у вас одна машина с несколькими GPU, точнее работает Лайфхак 8, а видеопамять как прежде не суммируется.

Лайфхак 18

Использование torch.jit для компиляции моделей может значительно ускорить их выполнение или перекомпеляция в onnx. Вы можете применять torch.jit.trace() или torch.jit.script() для преобразования модели в оптимизированный формат, который работает быстрее, особенно при повторных вызовах. Это особенно полезно, если вы часто вызываете одну и ту же модель для разных задач.

import torch

# Example of using torch.jit to trace a model
model = ...  # model
example_input = ...  # sample input suitable for your model
traced_model = torch.jit.trace(model, example_input)

# Now you can use traced_model instead of the original model
output = traced_model(example_input)

Лайфхак 19

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

import torch
from torch.profiler import profile, record_function

with profile(profile_memory=True) as prof:
    with record_function("model_inference"):
        output = model(input_data)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

И вот мы подошли к завершению нашей статьи с 19 лайфхаками! Хотя это и не круглое число, я чувствую, что не хватает ещё одного. Поэтому, пожалуйста, делитесь в комментариях вашим 20-м лайфхаком, чтобы сделать этот список полным.

Лирическое завершение

У меня есть мечта — увидеть 4096 звёзд на GitHub за мой проект. Я верю, что в топе GitHub должно быть больше проектов от русскоязычных разработчиков, и ваша поддержка даёт мне силы и вдохновение продолжать. Она позволяет мне улучшать код, разрабатывать новые подходы и делиться опытом. Если вам понравился мой труд, поддержите проект — и я обязательно продолжу создавать полезные материалы и делиться новыми идеями. А ещё расскажите о своих проектах с нейросетями на GitHub 🖐 — в комментариях!