Полтора года назад я начал работу над проектом с открытым исходным кодом, который постепенно рос и развивался. Вдохновившись проектом AUTOMATIC1111, на тот момент только появившимся, я добавлял всё больше функционала и возможностей. Сегодня мой проект включает более 50 нейронных сетей, каждая из которых выполняет свою уникальную задачу. В этой статье я делюсь практическими лайфхаками и выводами, которые помогли мне на этом пути. Надеюсь, что они будут полезны и вам.
Проект ориентирован на создание и редактирование видео, изображений и аудио с применением нейронных сетей. Часто разные методы могут выполнять схожие задачи. Так как я интегрировал решения с открытым исходным кодом, оптимизировал их и добавлял новый функционал, ключевой задачей стало обеспечение единства методов. Например, такие функции, как замена лица, синхронизация губ и анимация портретов, требуют распознавания лица. В моем проекте за эту задачу отвечает одна модель, а не несколько разных методов, как в исходных решениях. Поэтому все 50+ моделей распределены так, что каждая отвечает за своё уникальное направление, без дублирования.
В процессе разработки я принципиально отказался от TensorFlow и связанных с ним решений, сосредоточившись исключительно на PyTorch и ONNX Runtime.
Для тех, кто хочет детальнее ознакомиться с функционалом и узнать, какие именно нейронные сети я использовал, предлагаю несколько ссылок: плейлист на YouTube, где можно проследить, как проект развивался и совершенствовался, а также короткое видео, созданное с помощью моей программы — для тех, у кого нет доступа к YouTube.
Функционал каждой модели разнообразен и сложен: от генерации изображений и видео до распознавания лиц, сегментации и многого другого. В проекте нет простых решений, и каждая модель выполняет свою уникальную задачу.
Первое, с чем я столкнулся, и что стало для меня удивлением: нельзя загрузить одну модель в видеопамять и использовать её одновременно для нескольких задач. Необходимо, чтобы каждая модель загружалась под свою отдельную задачу. Следовательно, это станет основой для дальнейших лайфхаков.
Очередь. Мое приложение основано на Flask, поэтому пользователь не ожидает окончания обработки и может запускать сколько угодно задач, тем самым загружая память. В результате я искусственно создаю задержку между задачами с случайным значением, чтобы избежать одновременного запуска двух и более задач. Это связано с Лайфхаком 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
Вместе с отложенным запуском я использую проверки на самую распространенную ошибку: “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`, всё должно быть немного иначе.
Backend моей программы состоит из модулей, которые классифицируются по следующим признакам: изменение видео или изображения, генерация видео и изображений, изменение аудио — т.е. по свойству модели. И также по признаку: модель обрабатывает задачи для frontend или backend, т.е. результат работы модели необходимо вернуть мгновенно пользователю (сегментация, txt2img и img2img) или как выполненную крупную задачу. Мы не говорим про модели, которые работают на frontend, используя:
await ort.InferenceSession.create(MODEL_DIR).then(console.log("Model loaded"))
Следовательно, мне необходимо загружать модели для быстрого возврата ответа в память и держать их там, не позволяя разным пользователям одновременно использовать одну модель (Лайфхак 1) и не использовать их для задач с долгой обработкой, чтобы не нарушить Лайфхак 1.
Модели для длительной обработки иногда бывают очень требовательными, и в зависимости от видеопамяти, такая модель может полностью её использовать. В плане оптимизации очень невыгодно каждый раз загружать и выгружать такие модели, хотя иногда это, к сожалению, приходится делать. Часто с такими моделями используются микро модели, которые занимают в памяти немного места, но их загрузка и выгрузка требует времени. При запуске задач мы группируем их по методам длительной обработки, и задачи из одной группы обрабатываются на маленьких моделях, создавая очередь перед загрузкой в одну большую модель. Помните Лайфхаки 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
После выполнения каждой задачи очищайте память и удаляйте переменные и модели, если они больше не требуются.
del ...
Модели можно загружать по слоям на 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',
}
Не забывайте использовать 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.")
Такие подходы уменьшают количество используемой памяти, но увеличивают время обработки.
Не храним кадры в памяти. На самом деле – это палка о двух концах. Если вам нужно быстро получить результат на мощной машине с ограничениями по разрешению и длительности контента, то хранение в памяти может быть выгодным. Однако пользователи моего проекта запускают его на слабых устройствах с часовыми видео в высоком разрешении. Поэтому я переписал все методы для работы с текущим кадром и значениями, сохраняя их на жестком диске. Обращение к этим данным по мере необходимости позволяет избежать множества ограничений на устройства. В списке я храню только ссылки на файлы, что делает процесс более эффективным. Дополнительно можно использовать генераторы или чанки для обработки только текущих значений, подобное я делаю в некоторых модулях, например при замене лиц
Разрешение кадра. В зависимости от модели иногда приходится изменять размер кадра до пределов, которые может обработать устройство пользователя, а затем восстанавливать его размер обычным изменением размера или более продвинутым upscale.
Модели не бывают асинхронными? Это не является утверждением, так как мир искусственного интеллекта, постоянно меняется и это только мой опыт. Я обнаружил, что не получаю значительных выигрышей от использования асинхронных методов, за исключением отдельных операций обработки данных, которые не связаны напрямую с моделью, а также requests для загрузки и проверки актуальности модели. Модели работают синхронно.
Давайте поговорим о совместимости версий библиотек, особенно таких, как 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
Чтобы убедиться, что 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
Что делать, если некоторые модели работают только со старыми библиотеками, а другие — только с новыми? Я столкнулся с такой проблемой в 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
Таким образом, вы импортируете измененные методы для тех, которые исчезли в новых версиях. Это позволяет обеспечить совместимость между различными библиотеками и моделями.
Обращайте внимание на предупреждения (Warning). Всегда следите за сообщениями типа Warning, в которых говорится о предстоящих изменениях в новых версиях библиотек. Ищите соответствующие строки кода в вашем проекте и добавляйте или изменяйте необходимые параметры. Это поможет избежать накопления несоответствий при обновлении до новых версий.
Управление GPU в кластере. Если вы используете кластер из нескольких машин, помните, что вы не можете суммировать видеопамять от разных GPU. Однако, если видеокарты находятся в локальной сети, вы можете использовать управление GPU из одного контроллера. Для этого существуют библиотеки, такие как Ray. Обратите внимание, что суммирование видеопамяти не работает, за исключением случаев, когда у вас одна машина с несколькими GPU, точнее работает Лайфхак 8, а видеопамять как прежде не суммируется.
Использование 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)
Используйте инструменты профилирования, такие как 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 🖐 — в комментариях!