Полтора года назад я начал работу над проектом с открытым исходным кодом, который постепенно рос и развивался. Вдохновившись проектом 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, в которых говорится о предстоящих изменениях в новых версиях библиотек. Ищите соответствующие строки кода в вашем проекте и добавляйте или изменяйте необходимые параметры. Это поможет избежать накопления несоответствий при обновлении до новых версий.
Лайфхак 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 ? — в комментариях!