Всем привет! Попалась мне как-то на глаза публикация про квантизацию локальных моделей Whisper, данные модели я применяю для расшифровки аудио и видеороликов в рамках проекта Генератор субтитров (подробнее можно почитать в этой статье).
Так вот, в упомянутой публикации про квантизацию демонстрируется отличный пример того как эту самую квантизацию выполнить, но мне было лень встраивать это решение в API-сервер, хотелось обойтись незначительными правками, поэтому пришлось идти своим путём :)
Введение
Существует множество решений для преобразования аудио в текст, есть как базовые вещи вроде модели speech2text или SOVA ASR заканчивая большими API сервисами наподобии Yandex SpeehKit или Сбер SaluteSpeech.
Когда я только начинал проект Генератор субтитров мне пришлось попробовать многие решения, в том числе и упомянутые в предыдущем абзаце, и больше всего меня расстраивало то что качество распознавания у них либо оставляло желать лучшего, либо же качество было приемлемым, но системы не позволяли возвращать текст в формате субтитров (и требовалась доп.обвязка для преобразования), либо же их работа стоила каких-то не малых денег (учитывая то какой объём текста я планировал через Генератор субтитров расшифровать), что для хобби проекта неприемлемо.
Поэтому пришлось искать дальше, нужна была модель которую было возможно запустить локально, желательно через Docker и чтобы она могла работать в формате API-сервера и после продолжительных поисков я нашёл решение, которое полностью соответствовало моим потребностям, а именно проект Whisper от "Open"AI. Модели данного проекта относительно небольшие, они могут генерировать ответ в формате JSON, TXT или SRT, могут определять язык и даже делать что-то типа перевода (который получался не очень хорошим, но приемлемым), помимо этого стоит обратить внимание на скорость их работы, они действительно очень шустрые, не чета даже продвинутым платным облачным решениями.
Модели Whisper
На данный момент существует несколько разновидностей моделей Whisper. Веса моделей имеют тип данных float32 (8 бит - 1 байт, 32 бита - 4 байта), что говорит нам о том, что возможно несколько разных путей оптимизации всех перечисленных моделей, но об этом позже, для начала посчитаем сколько VRAM надо только для запуска в "базовой комплектации".
- tiny ~39M
Имеет 39 миллионов атрибутов, используем формулу из моей предыдущей публикации и рассчитаем сколько это в байтах.
39 000 000 * 4 = 156 000 000 ≈ 148.7Мб
- base ~74M
74 000 000 * 4 = 296 000 000 ≈ 282.2Мб
- small ~244M
244 000 000 * 4 = 976 000 000 ≈ 930.7Мб
- medium ~769M
769 000 000 * 4 = 3 076 000 000 ≈ 2933.5Мб
1 550 000 000 * 4 = 6 200 000 000 ≈ 5912.7Мб
Иными словами столько памяти минимум необходимо для запуска упомянутых ранее моделей, но помимо этого стоит учитывать также и то, что модели резервируют часть доступной памяти под инференс. То есть можете смело умножать полученные числа на два чтобы узнать реальное количество необходимой памяти.
Причины оптимизации
Почти все перечисленные модели имеют две версии: мультиязыковая и только английский, почти потому что large версии поставляются только в мультиязыковом формате.
Изучив всё это и протестировав качество работы моделей я пришёл к выводу, что для моего хобби проекта логичнее всего использовать либо модель medium либо одну из large, после дополнительных тестов пришёл к выводу, что только large соответствует моим ожиданиям качества.
Однако, как не сложно догадаться, для запуска large требуется порядочное количество ресурсов (12Гб видеопамяти в частности), и если на момент публичного запуска Генератора субтитров это было не столь критично, то спустя полгода и обзаведясь постоянной аудиторией я начал искать способы как на одной 16Гб видеокарте запустить две или более моделей.
Первая мысль была такая: раз на одну видеокарту не влезает более чем одна large модель, то возможно я смогу запустить вторую или третью на процессоре, однако, замеры скорости работы модели на процессоре показали, что среднее время расшифровки увеличилось десятикратно, а с учётом того что модель large рашифровывает примерно со скоростью x1 (в смысле расшифровка занимает столько же времени сколько длится аудиоролик) можно сделать предположение, что на процессоре работа будет выполняться со скоростью x0.1. Иными словами, если пользователь моего бота захочет запустить на процессорах расшифровку 4х часового аудиоролика, то Whisper будет занят этой работой примерно 16 часов, что меня не очень радовало.
Про Whisper ASR Web-service
Для запуска моделей Whisper существует несколько готовых решений, самый простой и удобный из них это Docker-контейнер whisper-asr-webservice, это пожалуй самый простой и функциональный проект из всех с которыми мне довелось работать, но он имеет ряд недостатков, он (на момент публикации поста) не позволяет выполнять инициализацию моделей в усечённом режиме (квантизация до 4/8 бит "на лету"), но может запускать модели в режиме float16.
Если запустить данный проект и зайти на порт 9000 то перед нами возникнет OpenAPI/Swagger веб-интерфейс, через который можно вручную выполнить отправку аудиофайла на вход модели и спустя некоторое время получим результат.
Однако, перед этим разберёмся с тем как запустить приложение.
Запуск приложения
Существует два варианта запуска модели: вручную через gunicorn и при помощи Docker-конетейнера. Разберём оба этих варианта, а также то какими переменными окружения пользоваться и как их передавать.
На данный момент возможно управлять приложением через следующие переменные окружения:
- ASR_ENGINE (default: openai_whisper)
- ASR_MODEL (default: base)
Начнём с запуска на хостовой операционке.
Запуск на хосте
Для запуска нам понадобятся следующие программы:
- Python 3.10
- Python VirtualEnv
Далее опционально, если будет необходимо запускать модель на видеокарте:
- Nvidia Driver (535.xx или выше)
- Nvidia CUDA (11.7 или выше)
Клонируем репозиторий:
git clone https://github.com/ahmetoner/whisper-asr-webservice.git
cd whisper-asr-webservice
Далее выполним инициализацию виртуального окружения:
python3.10 -m venv venv
Далее изменим контекст оболочки на него:
source venv/bin/activate
Теперь установим приложение poetry, оно потребуется для того чтобы управлять пакетами и зависимостями, ближайший аналог это npm или composer.
pip install poetry
Далее нам при помощи этой утилиты необходимо будет установить ряд пакетов, зафиксированных в poetry.lock
poetry install
Потребуется некоторое время для того чтобы приложение скачало все необходимые зависимости, далее перед нами вилка:
- Если у вас имеется свежая видеокарта от Nvidia то необходимо установить пакет torch оптимизированный для работы с GPU
pip install torch==1.13.0+cu117 -f https://download.pytorch.org/whl/torch
- Однако, если вы не планируете использовать видеокарту и хотите выполнять запуск модели на процессоре, то ничего более выполнять не требуется.
Теперь мы можем запустить приложение следующей командой:
gunicorn --bind 0.0.0.0:9000 --workers 1 --timeout 0 app.webservice:app -k uvicorn.workers.UvicornWorker
Если выполнить команду выше, то приложение будет инициализировано в режиме настроек по умолчанию, однако, если вы хотите изменить движок или тип модели необходимо передать переменные окружения в момент запуска gunicorn.
AWS_ENGINE=faster_whisper ASR_MODEL=large-v2 gunicorn --bind 0.0.0.0:9000 --workers 1 --timeout 0 app.webservice:app -k uvicorn.workers.UvicornWorker
Теперь мы можем открыть браузер и перейти по адресу http://localhsot:9000 и там увидеть OpenAPI/Swagger.
Запуск через Docker
Теперь поговорим про запуск при помощи Docker Engine (кстати, у меня есть две большие статьи про Docker, первая про основы, вторая про хитрости, рекомендую ознакомиться).
Для этого нам потребуется:
- Docker Engine
- Docker Compose
Далее опционально, если будет необходимо запускать модель на видеокарте:
- Nvidia Driver (535.xx или выше)
- Nvidia CUDA (11.7 или выше)
Выполним клонирование проекта Whisper Asr Web-service, но переходить в директорию проекта не будем:
git clone https://github.com/ahmetoner/whisper-asr-webservice.git
Теперь создадим рядом с этой директорией ещё одну, назовём её скажем whisper-tests.
mkdir whisper-tests
cd whisper-tests
В ней создадим файл docker-compose.yml и в зависимости от того настроен ли у вас Nvidia Docker Runtime или нет опишем файл.
Если вы хотите использовать уже готовый Docker-контейнер, собранный авторами проекта, то можете убрать комментарий с поля image, и добавить комментарий на блоке build. Однако, нам для дальнейших экспериментов более всего понадобится вариант собираемого локально контейнера, хотя полагаю после того как мой Pull Request будет смерджен то можно будет обойтись и без этого.
Если у вас нет настроенного Nvidia Docker Runtime то надо будет убрать блок deploy и вместо Dockerfile.gpu написать просто Dockerfile.
Выполним сборку:
docker-compose build
Запустим контейнер:
docker-compose up -d
Если зайти по адресу http://localhost:9000 то мы окажемся на странице сгенерированной при помощи OpenAPI/Swagger, на ней можно выполнять API запросы.
Оптимизируем-с
Сразу хочу заметить, что в дальнейшем я буду использовать только Docker версию приложения и используя при этом Docker Nvidia Runtime, так что создадим следующий docker-compose.yml конфиг, чтобы иметь точку отсчёта относительно которой мы будем вносить изменения.
ASR_ENGINE можно не указывать, она по умолчанию равна openai_whisper.
Ссылка на пример конфига:
В данном конфиге мы используем nvidia-docker-runtime, образ onerahmet/openai-whisper-asr-webservice:latest-gpu оптимизированный для работы на видеокартах, модель whisper-large-v2, запуск осуществляется в режиме openai_whisper (без квантизации), а также открываем порт 9000 у веб-интерфейса.
Теперь запустим контейнер и сделаем фактические замеры VRAM в режиме простоя:
Как видно она сразу заняла 9.9Гб VRAM.
float16
Самый простой способ сократить объём памяти, необходимой модели large, это просто запустить её в режиме half (то есть float16), данное решение позволит вдвое сократить объём необходимой памяти, а качество расшифровки останется на примерно том же уровне, впрочем чтобы не быть голословным мы с вами это протестируем чуть позже.
Теперь попробуем запустить проект в режиме faster_whisper (который загружает модель в формате float16), автор API-сервера предусмотрел данную возможность, поэтому нам нужно лишь немного подкорректировать docker-compose.yml следующими образом:
Тут была изменена только переменная ASR_ENGINE.
Посмотрим на фактически использованные ресурсы:
В режиме faster_whisper видно, что модели необходимо 6.6Гб VRAM, неплохо, но мы можем сделать ещё лучше.
We need to go deeper
Теперь попробуем разобраться с тем как происходит инициализация модели в режиме float16 когда мы используем движок faster_whisper. Найдём исходники проекта whisper-asr-webservice, в нём зайдём в директорию app/faster_whisper/ и в ней обнаружим файл core.py.
Там на 14й строчке происходит вызов функции model_converter:
Посмотрим, что она делает:
На скриншоте видно, что происходит запуск класса TransformerConvert, после чего происходит запуск метода convert. Тут стоит обратить внимание на параметры, в частности на уровень квантизации, в текущей реализации там захардкожено значение float16, но если судить по документации к методу convert мы можем заменить его на некоторые другие значение, вот список возможным вариантов:
- int8
- int8_float16
- int16
- float16
- int4 (в документации он не упоминается, но исключения не выбрасывает, так что возможно работает)
Давайте слегка её модифицируем, добавив возможность указывать уровень квантизации вручную:
Тут видно, что вместо хардкода я сделал параметр quantization со значением по умолчанию, равным float16.
Теперь вернёмся в файл core.py и обновим его таким образом, чтобы не сломать существующую логику, а выбор режима квантизации сделаем опциональныи, но при этом настраиваемым.
Начнём описание сверху, первым делом я обозначил массив содержащий в себе логику, которая выполняет маппинг, не знаю зачем автор проекта выполняет квантизацию до 16 бит, а потом инициализирует модель в режиме 32 бита, но возможно в этом есть какая-то неведомая мне причина, поэтому я постарался сохранить эту структуру для остальных уровней квантизации через маппинг.
Уважаемые читатели!
Если у вас есть объяснение причины по которой автор проекта Whisper Asr Web-service реализовал подобную логику, то напишите пожалуйста об этом в комментариях к посту у меня в Telegram-канале.
На всякий случай я создал Issue на странице проекта, возможно автор сможет дать соответствующие комментарии.
Если что-то новое узнаю, то обновлю публикацию.
Далее мы считываем значение переменной окружения ASR_QUANTIZATION (если она не задана, то имеет значение float16, нужно это для обратной совместимости), следом идёт чтение переменно ASR_MODEL равное base по умолчанию, после чего мы обозначает путь до домашней директории пользователя, а в ней до директории куда проект будет сохранять квантизированные модели.
На следующем шаге мы изменяем путь сохранения моделей, выглядеть он будет так: <модель>_<квантизация>, для float16 подобная задача не выполняется (для обратной совместимости).
Далее мы вызывает функцию конвертации модели передавая уровень квантизации третьим параметром, после чего используя MAPPING определяем какой compute_type необходимо передать на вход объекта модели Whisper.
Вот в принципе и всё, теперь мы можем пересобрать Docker-контейнер, запустить четыре версии модели (float32, float16, int8, int4) и выполнить замеры скорости расшифровки аудиороликов, объём памяти, который при этом используется и естественно время необходимое для расшифровки.
Паноптикум квантизаций
На всякий случай упомяну, что у меня имеется сервер с видеокартой RTX 4090, которая обладает 24Гб оперативной памяти, что позволяет мне без особых сложностей запустить два экземпляра модели Whisper large v2 и не беспокоиться о том, что память может внезапно иссякнуть.
Поэтому я могу запустить сразу четыре экземпляра модели, для этого создам специальный docker-compose.yml файл, который будет содержать в себе схожие конфигурации, отличающиеся только значением переменной ASR_QUANTIZATION (либо же она будет отсутствовать для режима float32).
Для начала укажем версию и перечислим переиспользуемые шаблоны YAML (об этом я подробно рассказывал в публикации про хитрости Docker), они будут повторяться во всех контейнерах.
Далее опишем блок services и перечислим в нём четыре сервиса:
- whisper01 будет работать в режиме float32, порт 9001
- whisper02 - float16, порт 9002
- whisper03 - int8, порт 9003
- whisper04 - int4, порт 9004
Полный код данного docker-compose.yml вы сможете найти тут.
Запускаем контейнеры поочерёдно и делаем замеры VRAM которая была задействована для этого, значения float32 (9.9Гб) и float16 (6.6Гб) нам уже известны, поэтому пропустим их и остановимся на int8 и int4.
docker-compose up -d whisper03
Посмотрим сколько памяти требуется:
Тут видно, что модель large-v2 в режиме квантизации int8 необходимо всего лишь 3.4Гб оперативной памяти, теперь повторим процедуру запуска для режима int4.
docker-compose up -d whisper04
В режиме int4 модель large-v2 требует лишь 1.9Гб VRAM.
Что там с качеством?
Чтобы проверить насколько хорошо работает каждая из версий модели необходимо провести тестирование, действовать будем в соответствии со следующими планом:
- Подготовка моделей
У нас уже готовы и запущены 4 версии модели large-v2. - Выбор эталонных аудиороликов
Нам потребуется 10 аудиороликов: 5 на русском языке и 5 на английском. Продолжительность каждого ролика должна быть не очень большой (5-30 минут), дабы тесты не затянулись на неопределённое время. Помимо этого заранее выполним эталонную расшифровка этих аудиороликов при помощи модели large-v2 в режиме float32. - Расшифровка аудиороликов
Каждая квантизированная версия модели будет расшифровывать все 10 аудиороликов, для этого напишем специальный скрипт, который будет работать с API эндпоинтом /asr. - Оценка точности
Сравниваем результат расшифровки с эталонной и вычисляем точность работы каждой модели (относительно эталона). Для этого мы воспользуемся алгоритмом вычисления расстояний Ливенштейна. - Замер времени
Для каждого аудиоролика мы также измеряем, сколько времени модель потратила на его расшифровку. - Результаты
Все полученные данные собираем в одну сводную таблицу. В этой таблице для каждой модели будет указана точность расшифровки и среднее время обработки аудиоролика.
Таким образом, мы сможем оценить, какая модель работает быстрее и точнее, но на самом деле меня интересует есть ли отличие в точности расшифровки в принципе, потому как моя гипотеза в том что их либо вообще нет, либо отличия в точности незначительны.
Скрипт автоматизации расшифровки
Набросаем небольшой скриптик, который будет отправлять mp3 файлы на нужный порт сервиса (доступного на localhost).
Исходники тут.
Данный скрипт поочерёдно отправляет файлы с названиями от 1 до 10 на эндпоинт /asr, перед каждым шагом запускает подсчёт времени, после чего ожидает ответа Whisper, затем сохраняет файл на диск и переходит к следующему файлу. Как видно ничего сложного.
Из примечательного пожалуй только переменная PORT которая позволяет определять нам на какой сервер происходит отправка файлов, помимо этого значение переменной используется для сохранения txt расшифровок.
float32 (эталон)
А вот эталонные значения для 10 видеороликов:
На этой таблице можно заметить, что в режиме float32 видеокарта RTX 4090 производит расшифровку со скорость примерно x10, то есть в десять раз быстрее чем воспроизводится аудиодорожка. А ещё что для выполнения перечисленного количества тестов потребовалось чуть меньше 11Гб VRAM.
Среднее время расшифровки 10 видеороликов составило 132.5 секунд.
float16
Теперь проведём тесты в режиме float16, но перед этим добавим в скрипт логику расчёта похожести двух текстов, существует множество замечательных алгоритмов, например расстояния Левенштейна, косинусное сходство, Jaccard-индекс, но я решил использовать расстояния Левенштейна, так как мне это показалось самым простым способом.
Слегка модифицируем скрипт:
И выполним тестирование, получилась следующая таблица:
В режиме float16 наблюдается незначительная потеря точности (в сравнении с эталоном), ошибка в среднем равна 4%.
Среднее время расшифровки 10 видеороликов составило 77.5 секунд, иными словами время расшифровки сократилось примерно на 42%.
Необходимый объём VRAM при этом сократился до 8.4Гб, что уже неплохо, двигаемся дальше.
int8
Поправим номер порта на 9003 и запустим скрипт по новой, вот таблица с результатами:
Видно что оперативки теперь надо чуть меньше 5Гб, среднее время расшифровки (по сравнению с эталоном) сократилось на 75% при этом точность расшифровки такая же как и у float16.
int4
Теперь натравим скрипт на порт 9004 и посмотрим что получилось:
Количество необходимой памяти оказалось ниже чем у int8, но не в два раза, как я ожидал, помимо этого любо я что-то не так сделал, либо TransformerConvert не поддерживает int4 (странно что не было эксепшена), но время расшифровки аудиодорожки немного увеличислось, оно разумеется по прежнему чуть ли не в два раза меньше чем на float32, но немного больше чем в int8 режиме.
Заключение
Неожиданно большая публикация получилась, изначально я планировал просто сделать небольшой перевод, но автор оригинальной публикации много моментов недоговорил, поэтому мне захотелось расширить эту тему, и конечно же провести доскональные тесты.
Для дочитавших до конца предлагаю ссылки на модифицированный мною проект whisper-ars-webservice и набор конфигов и скриптов которые я использовал в процесс выполнения тестов, там же будут результаты тестов.
Не забудьте поставить лайк, подписаться на канал, а также приглашаю ко мне на Telegram-канал на котором я публикую свои размышления, рассказываю о разных интересных новостях, а также делюсь своими наработками и опытом.
К тому же, если вы хотите поддержать мои усилия и вклад в развитие общества знаний, вы можете сделать пожертвование на CloudTips. Ваша поддержка поможет мне продолжать свою работу и делиться новыми открытиями с вами.
До встречи в следующей публикации!
