Мы уже писали о том, что предложили новую модель квантования нейронных сетей, позволяющую ускорить их на 40% на центральных процессорах, а также о том, как она устроена вот тут.

Сегодня мы расскажем о том, как мы в Smart Engines обучали 4.6-битные сети. Основная проблема обучения квантованных моделей в том, что градиентные методы модифицируют веса непрерывным образом, а у квантованных сетей они дискретные. 

Кроме того, распространен сценарий, когда сначала обучается вещественная сеть, а затем ее хочется отквантовать, чтобы ускорить систему. Часто обучающие данные при этом уже недоступны или доступны не в полном объеме.

Поэтому методы квантования делятся на 2 большие группы:

• Post Training Quantization (PTQ): когда мы квантуем уже обученную сеть. Как правило, они минимизируют ошибку между выходами каждого слоя вещественной и квантованной сетей. Такие методы хорошо работают на больших моделях, в которых есть избыточные веса, а в случае компактных сетей для центральных процессоров отквантоваться без снижения точности обычно не получается.

• Quantization Aware Training (QAT): когда мы уже во время обучения применяем методы, повышающие качество квантованной модели.

В Smart Engines все модели уже компактные, а обучающие данные никуда не деваются, так что нам в первую очередь интересны QAT методы или их комбинации с PTQ методами. Расскажем о них поподробнее.

Метод обучения

Самый очевидный способ обучить квантованную сеть – игнорировать квантование во время обратного прохода по сети, а модифицированные веса округлить. Тогда прямой проход будет уже квантованным и соответствовать “рабочему” режиму сети. Так работает Straight Through Estimator (STE), один из старейших и популярных методов обучения квантованных сетей. На удивление, он неплохо работает, если взять уже обученную 8-битную модель и ее постепенно квантовать, то есть 256 значений для сети оказывается не так уж и мало. Однако округление на каждой итерации делает процесс обучения “шумным” и может привести к невозможности обучить сеть вообще, что и наблюдается в случае меньших разрядностей.

Поэтому даже для 8-битных моделей есть методы получше, например, инкрементное обучение.

Идея инкрементного обучения или понейронной конвертации заключается в том, что мы постепенно преобразуем часть весов каждого слоя в квантованный вид. Уже квантованная часть сети немного дообучается (например, с помощью STE) и “замораживается” (веса фиксируются и больше не меняются). Вещественная часть сети дообучается, чтобы адекватно обрабатывать изменившиеся входные значения от квантованной части. Этот подход прост, эффективен в отношении 4- и 8-битных сетей, но требует времени.

AdaQuant – PTQ метод, работающий послойно. Его идея которого заключается в том, что мы используем небольшой калибровочный датасет, чтобы определить динамические диапазоны активаций, а также коэффициент масштабирования и точку нуля так, чтобы минимизировать среднеквадратичную ошибку между выходом квантованного и вещественного слоев.

В данной работе мы использовали комбинацию этих подходов. Сначала параметры каждого слоя отквантовали с помощью Adaquant, а затем применили инкрементное обучение, а именно прошли от начала к концу и еще немного дообучили каждый из квантованных слоев (остальные “замораживались”) с помощью стохастического градиентного спуска с моментом 0.9 и скоростью обучения (learning rate) .

Далее мы рассмотрели три задачи компьютерного зрения, решаемые с помощью разных нейросетевых архитектур, и применили в них 4.6-битные сети. Кстати, наш код экспериментов есть на GitHub: https://github.com/SmartEngines/QNN_training_4.6bit

CIFAR-10

Это давно набившая оскомину небольшая выборка с цветными изображениями объектов 10 классов размера 32 на 32 (см. примеры на рис. 1). В качестве аугментации использовали отражения вдоль вертикальной оси, вырезание случайных регионов со сдвигами и случайные повороты на +-9 градусов с дополнением изображений на 4 пикселя с краев.

На ней мы обучили несколько простых сверточных моделей. Их архитектуры приведены в таблице 1. 

Обозначения:

• conv(c, f, k, [p]) – сверточный слой с c-канальным входом, f фильтрами размера k х k и паддингом p (в обоих направлениях, по умолчанию 0),

• pool(n) – 2D max-pooling с окном n на n,

• bn – batch normalization,

• HardTanh – активация HardTanh(x) = min(1, max(-1, x)),

• relu6 – активация ReLU6(x) = min(max(0, x),6),

• tanh – гиперболический тангенс,

• fc(n) – полносвязный слой с n выходами.

Таблица 1. Архитектуры сверточных сетей для CIFAR-10. 

Сначала мы обучили вещественные версии этих сетей 250 эпох с оптимизатором AdamW с параметрами по умолчанию, кроме коэффициента убывания весов (weight decay), который был установлен в , и скорости обучения (learning rate), которая была задана как и уменьшалась вдвое каждые 50 эпох. Размер батча был 100.

Дальше слои батч нормализации были интегрированы в свертки и мы приступили к квантованию. Первый и последний слой мы не квантовали, так как это улучшает результаты обучения и практически не снижает вычислительную эффективность. Результаты приведены в таблице 2. Напомним, что у нашей схемы квантования есть параметры и , которые обозначают количество дискретов для активаций и весов соответственно, мы перебрали все возможные их комбинации. Эксперименты повторили 5 раз с разной начальной инициализацией сети и посчитали среднюю точность и погрешность.

Таблица 2. Точность классификации на выборке CIFAR-10

Можно видеть, что 8-битные модели демонстрируют почти ту же точность, что и вещественные, а вот 4-битные им сильно уступают. Для 4.6-битных моделей оказалось, что лучше всего распределять дискреты между активациями и весами равномерно:  наилучшие результаты показывают схемы квантования от (15, 37) до (31, 17). Отметим, что точность классификации в этом диапазоне заметно лучше, чем у 4-битных моделей. Тем не менее, 4.6-битные сети все же немного уступили 8-битным и вещественным моделям, что говорит о том, что метод обучения можно придумать и получше.

ImageNet

Это уже стандартная выборка для оценки точности классифицирующих моделей (см. примеры на рис. 2). 

Здесь мы использовали ResNet’ы: ResNet-18 (11.7М весов) и ResNet-34 (21.8М весов), предобученные модели из PyTorch's Torchvision. Все активации ReLU заменили на ReLU6, батч нормализацию интегрировали в свертки и отквантовали, но первый и последний слои также не трогали. 

Калибровочная выборка была меньше: 25 батчей по 64 изображения для ResNet-18 и 5 батчей по 64 изображения для ResNet-34.

Точности классификации top-1 и top-5 для выборки ImageNet показаны в таблице 3. Лучшая пара параметров для 4.6-битных моделей оказалась (, ) = (23, 23). Также как и на CIFAR-10 8-битные модели имеют сравнимое качество с вещественными, а 4-битные сильно хуже. 4.6-битные сети оказались почти ровно посередине. Это ухудшение не такое большое и с ним уже можно работать!

Таблица 3. Точность классификации на ImageNet.

TCIA

Это выборка с изображениями МРТ головного мозга, изображения RGB, размер каждого 256 на 256. На этих изображениях ставится задача сегментации: нужно найти и отметить аномалии. Эта задача прекрасно  решается с помощью модели U-Net с 7.76М параметров. Мы взяли предобученную модель тут. Она принимает на вход изображение и выдает бинарную маску, которая отмечает аномалии. 

Также как и в ResNet’ах мы заменили активации ReLU на ReLU6 и отквантовали сеть. Размер калибровочной выборки был 5 батчей по 10 изображений. Для оценки качества сегментации использовались средние значения Dice и Intersection over Union (IoU) для тех изображений, где аномалии были. Также посчитали метрики качества бинарной классификации (есть аномалия или нет): accuracy, precision, recall, ошибку 1-го рода (false positive rate) и ошибку 2-го рода (false negative rate). Метрики приведены в таблице 4, а примеры работы на рис. 3.

Таблица 4. Качество сегментации моделей U-Net на выборке TCIA.

Можно видеть, что 8-битная сеть даже немного лучше вещественной: у нее чуть больше ошибка 1-го рода, но при этом выше качество сегментации. 4-битной сети откровенно плохо как по метрикам, так и визуально: добавился заметный шум. А 4.6-битная модель работает разумным образом и демонстрирует качество близкое к качеству 8-битной и вещественной моделей.

Обсуждение

Наши эксперименты ясно показали, что даже при использовании простейших методов обучения, не слишком отличающихся от подходов для 4- и 8-битных моделей, 4.6-битные сети работают заметно лучше 4-битных при практически такой же скорости работы. Это было ожидаемо, так как 4.6-битная схема квантования предлагает больше дискретов, чем 4-битная, а мы только проверили это теоретическое обстоятельство на практике. 

На сегодняшний день в нише быстрых моделей для центральных процессоров основную роль играют 8-битные модели за счет сочетания факторов: заметное ускорение и малые усилия по конвертации к 8-битному формату. Именно их наиболее актуально сравнивать с 4.6-битными. Здесь 4.6-битные сети однозначно выиграли по скорости, продемонстрировав ускорение на 30-40%. При использовании простых в имплементации методов обучения они несколько проиграли 8-битным по качеству, однако можно взять чуть более сложную 4.6-битную модель, которая даст сходное качество, но будет все еще работать быстрее исходной. Поэтому на практике мы получаем ускорение без снижения качества, а также активно работаем над усовершенствованием метода обучения 4.6-битных сетей.

На самом деле здесь есть много простора для фантазии:

• параметры , позволяют модифицировать схему квантования под разные задачи,

• никто не говорит, что они должны быть одинаковые в разных слоях модели,

• AdaQuant все-таки устарел, есть и другие методы, которые могут подойти лучше.

Поэтому мы считаем 4.6-битные сети крайне перспективными для практического использования и дальнейшего усовершенствования. Мы уже используем их в наших системах распознавания паспортов, распознавания документов, а также подали заявку на патентование этой технологии в РФ и США.