Оптимизация моделей в Tensorflow

Современные модели данных для машинного обучения, в основном, довольно тяжелые, поэтому работать с ними при ограниченных аппаратных ресурсах может быть проблематично. Уменьшить объем модели в TensorFlow и сократить расчеты без ущерба для качества обучения поможет оптимизация, причем для этого разработано немало методов. В статье познакомим вас с ними, но сначала о том, как устроены сами модели — это поможет легче понять, как их оптимизировать.

Модели обучения в TensorFlow: структура и стандартные методы хранения

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

Также особым образом устроено и хранение моделей. Например, граф вычислений и веса чаще всего хранятся отдельно: во всяком случае при использовании стандартных методов хранения. А их два: Checkpoint и SavedModel. Чекпоинты — это контрольные точки обучения, в которых сохраняются точные значения текущих параметров, но без вычислений, которые были определены моделью. SavedModel, напротив, подразумевает сохранение и таких вычислений. Этот метод реализован при помощи API tf.saved_model, который облегчает работу с моделью, добавляя поддержку различных фреймворков и CLI. Также SavedModel поддерживает ассеты, значительно расширяющие возможности работы с моделью. Так, вы можете добавлять различные списки: например, меток и стоп-слов.

Однако о методах хранения Checkpoint и SavedModel очень подробно написано в официальной документации TensorFlow: соответственно, здесь о Checkpoint, а здесь о SavedModel. Мы же рассмотрим те, которым в документации разработчиков уделено значительно меньше места, или они только упоминаются.

Современные методы оптимизации моделей в TensorFlow

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

Метод заморозки графа

Один из редких случаев, когда граф вместе с переменными и весами сохраняются в один файл. При этом удаляются все элементы графа, которые не требуются для непрерывной работы нейросети на конечных устройствах. Правда, платой за это является невозможность дальнейшего обучения модели стандартными методами. Тем не менее заморозка графа позволяет сокращать объем моделей в 3-4 раза.

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

graph = tf.get_default_graph()

input_graph_def = graph.as_graph_def()

Теперь выполняем собственно заморозку, что можно сделать по-разному. Часто задействуют метод freeze_graph, который подробно задокументирован на GitHub, с кодом и комментариями. Также можно преобразовать переменные в константы, используя следующую инструкцию:

output_graph_def = graph_util.convert_variables_to_constants(self.sess, input_graph_def, output_node_names)

И выполнить собственно заморозку:

with tf.io.gfile.GFile('graph.pb', 'wb') as f:

    f.write(output_graph_def.SerializeToString())

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

Метод выборочного сохранения

Также используется название «сохранение легкого Checkpoint». Благодаря модулю tf.train.Saver появляется возможность выборочно сохранять переменные и отключить сохранение метаграфа. Первое делается так:

saver = tf.train.Saver(var_list=tf.trainable_variables())

А вот как отключается метаграф:

ckpt_filepath = saver.save(self.sess, filepath, write_meta_graph=False)

Эти операции позволяют добиться тех же результатов в плане объема, что и первый метод. Но всегда нужно смотреть на те возможности для работы с моделью, которые остаются после применения того или иного метода. В данном случае бояться нечего, поскольку для непрерывной работы нейросети нам потребуются только веса, а чтобы выстроить граф, можно использовать дублирование кода, что делает ненужным и сохранение метаграфа.

Метод обнуления весов

Он также называется прунингом. Это один из тех методов, который позволяет несколько ужимать модели, но без заметных потерь в точности вычислений. Прунинг может быть выполнен разными способами:

• через оптимизацию графов при помощи Grappler;
• через Pruning API;
• при помощи Graph Transform Tool.

Все три способа увеличивают число нулевых значений в модели.

Оптимизация через Grappler — это «коробочное» решение, которое подробно описывается в официальной документации. Но эксперименты с настройками Grappler показывают, что размер модели после применения прунинга оказывается ненамного меньше.

Pruning API и Graph Transform Tool также не дают значительного выигрыша в объеме, если мы говорим о сверточных моделях. Дело в том, что в таких моделях почти всё пространство занимают эмбеддинги, прунинг к которым должен применяться отдельно. А в реализации этих методов по умолчанию прунингу подвергается лишь малая часть графа, что не сильно экономит место, а значит, выигрыш от работы с весами оказывается минимальным.

Метод квантизации

Этот способ подразумевает значительное сжатие коэффициентов весов, что может приводить к некоторому ухудшению общей точности работы модели. Но на практике основное преобразование quantize_weights не приводит к каким-либо отрицательным эффектам, зато позволяет сократить объем модели всё в те же 3-4 раза. Квантизацию сверточных моделей позволяет выполнять вот эта утилита. Если же ваши нейросети активно используются владельцами смартфонов и планшетов, присмотритесь к TensorFlow Lite, который поддерживает все современные методы квантизации.

Метод сокращения точности вычислений

Учитывая, что данные в моделях обучения TensorFlow представлены в векторном формате, даже 32-разрядные числа, используемые для их инициализации — это перебор, не говоря уже о 64-разрядных. Поэтому некоторое сокращение точности вычислений (до 16-разрядных значений после запятой) позволяет освободить определенное количество оперативной памяти (но не более 10% от используемой). Тем не менее пониженная точность данных может давать неожиданный эффект и приводить к существенному увеличению времени обучения, порой десятикратному. И с ходу трудно найти настройки, которые бы помогли достичь приемлемого времени обучения. Так что использование этого метода требует предварительного тестирования конкретной модели и в большинстве случаев не оправдывает себя из-за значительно увеличивающихся временных затрат на обучение.

Заключение

TensorFlow для машинного обучения — по-прежнему одно из лучших решений. И во многом это так благодаря гибкости библиотеки, которая позволяет работать с моделями самыми разными способами. Оптимизацию моделей тоже можно выполнять по-разному, причем какие-то методы работают лучше, а какие-то хуже для конкретной модели. И мы надеемся, что после прочтения этой статьи у вас осталось значительно меньше вопросов относительно оптимизации моделей машинного обучения.