Лайфхаки: 50+ нейронных сетей в одном проекте. Как работает и для каких задач?

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

Прежде чем начать, стоит отметить, что мой проект позволяет создавать и редактировать видео, изображения и аудио с помощью нейронных сетей. Порой различные методы выполняют схожие задачи. Так как я изучал различные решения с открытым исходным кодом, оптимизировал их и добавлял новый функционал, передо мной также стояла задача единства. Например, замена лица, синхронизация губ и оживление портрета используют распознавание лица. Чтобы обработать лицо, его необходимо сначала найти. Исходные решения использовали три разных метода распознавания лица, а в моем проекте за один функционал отвечает только одна модель. Поэтому каждая из более чем 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, где он требует старую версию torchaudio, в то время как для генерации видео необходимы новые версии 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, в которых говорится о предстоящих изменениях в новых версиях библиотек. Ищите соответствующие строки кода в вашем проекте и добавляйте или изменяйте необходимые параметры. Это поможет избежать накопления несоответствий при обновлении до новых версий.

Когда увидел сообщения в консоли

Лайфхак 17

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

Лайфхак 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-м лайфхаком, чтобы сделать этот список полным.

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

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