Для решения задачи снижения потребления памяти и повышения скорости расчета моделей нейронных сетей без существенной потери точности используется различные методы оптимизации.
Бывает так, что очень большая модель не помещается на видеоадаптере и требуется 250 ГБ оперативной памяти. В этой связи надо находить баланс, можно уменьшить размер модели в сто раз, при этом уменьшить точность всего на половину процентного пункта. Например, Bert можно сжать с 560 Мб до 2 Мб, почти без потери качества.
Рассмотрим три наиболее часто встречающихся метода оптимизации размера сети такие как дистилляция, квантизация и прунинг.
Основная идея дистилляции, это обучение маленькой модели (модели студента) с помощью предобученной большой модели (модели учителя). Пусть у нас есть предобученная модель «учитель», она выдаёт логиты, это последний слой до Softmax. И есть модель студента, только необученная, которая выдает логиты, такой же размерности. Далее мы логиты учителя и логиты студента, прогоняем через Softmax с температурой. Температура нужна для сглаживания распределения.
Softmax c температурой
Логиты модели-студента отправляются на Softmax как с температурой, так и без температуры.
При Softmax без температуры мы получаем классическую модель-студент обучения. И обычную функцию потерь. Эту функцию мы возьмем с коэффициентом альфа.
При Softmax с температурой мы сравниваем модель учителя и модель студента дивергенцией Кульбака – Лейблера, и этот loss берем с коэффициентом бетта. Потом складываем два loss как на схеме выше и берем от него backward. (обратное распределение потерь). Пример кода представлен ниже.
Иначе говоря, вместо классической функции потерь для обучения модели-студента, мы берем средневзвешенную функцию потерь с реальными данными и с теми что выдает модель учитель.
Квантизация
Идея квантизации предельно проста, все операции проводятся в целочисленных значениях. Чаще всего это накладывается на слой или на какую-то часть сети. Ряд слоёв работает в int8, благодаря чему он потребляет очень мало вычислений и памяти, последний слой, float32. Самый простой способ — это округление весов к определённым значениям.
Квантизовать можно до, после и во время обучения. Приведем пример простой квантизации для однослойной модели и четырех нейронов с четырьмя признаками
Кроме того, можно квантизовать уже дистиллированную модель, что ещё больше сэкономит время и вычислительные ресурсы, практически без серьёзной потери качества.
Прунинг
Прунинг нейронной сети это метод сжатия модели, путем удаления части параметров.
Прунинг, как и квантизация, может быть использован до, после и во время обучения на основе:
● Амплитуды весов, активаций, градиентов, гессианов
● Заданных правил, Байесовских подходов
● Реинициализации, дообучения
Самый простой способ прунинга нейронной сети есть под капотом PyTorch, и он реализуется всего одной строчкой кода.
В заключение хотелось бы сказать, что все три способа оптимизации можно комбинировать, тем самым получать результаты сопоставимые по качеству с огромными моделями, при этом затрачивая гораздо меньше ресурсов как временных, так и информационных.