Как ускорить инференс deep learning модели?

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

Дистилляция знаний (Knowledge distillation, KD)

Первое что приходит на ум, когда стоит задача ускорения модели и уменьшения занимаемой ею памяти - это дистилляция знаний. Кратко расскажу в чем заключается идея этого подхода: пусть у нас есть большая модель, которая хорошо решает определенную задачу, к примеру самый большой в мире ViT (у ViT H/14 ~633.5 миллиона параметров!!!) для классификации изображений и есть моделька поменьше, например ResNet50 (~25.6 миллиона параметров). Дистилляция знаний позволяет нам как бы перенести знания с большой модели на более компактную, в нашем примере с ViT на ResNet50. Большую модель обычно называют учителем, а меньшую модель называют учеником. Процесс дистилляции (который я рассмотрю подробнее ниже) можно проиллюстрировать так:

То есть в vanilla варианте дистилляции нам надо перенести предсказания с учителя к ученику в процессе обучения ученика (будем считать, что учитель уже обучен и хорошо справляется с задачей), но как это сделать? Очевидно, с помощью лосс-функции, то есть нужно добавить к обычному лоссу что-то, что будет приближать предсказания (то есть распределения предсказаний) учителя и ученика. Это что-то - KL-дивергенция. Простыми словами - это расстояние между распределениями. Тогда если наш лосс имел вид кросс-энтропии (т.е. минус сумма логарифмов вероятностей; в целом это мог быть любой классификационный лосс). То с добавлением KL-дивергенции, лосс принимает вид:

Где альфа - некоторый гиперпараметр, отвечающий за баланс между собственными предсказаниями ученика и "подражанием" своему сенсею. p_s - предсказания ученика, p_t - предсказания учителя.

Если же ученик и учитель имеют схожую структуру (к примеру два трансформера или два резнета), то мы можем помимо подражания в предсказаниях делать это еще и на уровне скрытых состояний сеток:

Легче всего это делать с помощью обычного MSE. Добавив это в лосс, получим:

Где h_t - скрытое состояние учителя, h_s - скрытое состояние ученика.

Скажу еще раз: это лишь краткое описание метода KD, существует довольно много различных модификаций по типу self-distillation и т. д., а также много подводных камней вышеописанного (к примеру скрытые состояния сеток могут быть разных размерностей).

Кроме уменьшения моделей, дистилляция знаний может применяться к примеру для увеличение робастности моделей, советую прочитать про это на paperswithcode: тык номер раз. Иллюстрации для этой части поста я взял с курса Тинькофф: тык номер два, очень советую начинающим в DL. В целом метод рабочий и обычно сеть-ученик очень незначительно уступает по качеству сети-учителю, а иногда даже и превосходит.

Еще ссылочки по дистилляции знаний:

1. тык номер три, классный пост по дистилляции.
2. тык номер четыре, вроде как одна из первы статей по KD на архиве.
3. тык номер пять, тоже классная статья.
4. тык номер шесть, мой репозиторий с реализацией дистилляции для разных пет-проектов.

Квантизация (quantization)

Тут у нас разговор будет довольно короткий - этот способ больше не про скорость инференса, а про вес модели. Обычно веса моделей хранятся в типе данных float32, то есть занимают 32 бита (4 байта), а если вспомнить сколько параметров у какого-нибудь GPT-3 (~175 миллиардов...), то сразу становится понятно, что такая тема мало где может поместиться, поэтому закономерная идея - уменьшить размер хранимых параметров до скажем 8 битов (1 байт) и представлять их типом int8. Естественно, после таких фокусов качество предсказаний упадет, но иногда легковесность модели бывает в разы важнее ее качества.

Использование "экзотических" типов данных

Сейчас уже не таких и экзотических, но возможно кто-то про них не слышал.

Google (или Microsoft...) создали два типа данных для ускорения тензорных вычислений на TPU: brain float 16 (bf16) и tensor flow 32 (tf32). В сравнении с обычными типами данных:

Как видно из иллюстрации, bf16 и tf32 занимают места намного меньше, чем float32, но имеют ту же точность.

Если вы используете TPU или видеокарту Nvidia, то все прекрасно, эти устройства поддерживают bf16, проверить это можно к примеру на torch, таким нехитрым образом:

>>> torch.cuda.is_bf16_supported()
True

tf32 поддерживается вроде как только для A100. Почитать про эти и другие типы данных можно тут: тык номер севен. Помимо обсуждаемого тут инференса моделей, эти типы данных повсеместно используются в обучении больших нейросетей, ускоряя его в несколько раз! Этот метод обучения называется mixed-precision, суть его состоит в том, что некоторые вычисления проводятся на менее точном типе данных, а некоторые на более точном. Почитать про это можно тут: тык номер восевен и тык номер девять.

Конвертирование модели во всякие TensorRT, ONNX...

Для работы модели на GPU, однозначно выбор падает на TensorRT, как это сделать для PyTorch можно посмотреть тут: тык номер десять. Это ускорит инференс аж в 6 раз! Кроме того, я слышал, что при конвертации модели PyTorch -> ONNX -> TensorRT происходит настоящая магия!

Если же вы хотите запускать модель на CPU, то можно воспользоваться библиотекой IPEX: тык номер одиннадцать, и вообще все это можно удобно обернуть на Hugging Face...

Заключение

Здесь перечислены далеко не все способы ускорения инференса deep learning модели, но начать можно с этого. Пост будет обновляться по ходу изучения этого вопроса, поэтому можете зайти сюда через недельку-две и найдете еще что-нибудь интересное (к примеру новые полезные ссылочки). Разумеется, приветствуются комментарии, критика и дополнительная информация по вопросу.

Всем спасибо за внимание!