Преобразование TensorFlow RNN в TensorFlow Lite

Обзор

TensorFlow Lite поддерживает преобразование моделей TensorFlow RNN в объединенные операции LSTM TensorFlow Lite. Объединенные операции существуют для максимизации производительности лежащих в их основе реализаций ядра, а также для обеспечения интерфейса более высокого уровня для определения сложных преобразований, таких как квантование.

Поскольку в TensorFlow существует множество вариантов API RNN, наш подход был двояким:

  1. Обеспечьте встроенную поддержку стандартных API TensorFlow RNN, таких как Keras LSTM. Это рекомендуемый вариант.
  2. Предоставьте интерфейс к инфраструктуре преобразования для пользовательских реализаций RNN , которые можно подключить и преобразовать в TensorFlow Lite. Мы предоставляем несколько готовых примеров такого преобразования с использованием интерфейсов LSTMCellSimple и LayerNormalizedLSTMCellSimple RNN от lingvo.

API конвертера

Эта функция является частью версии TensorFlow 2.3. Он также доступен через tf-nightly pip или через head.

Эта функция преобразования доступна при преобразовании в TensorFlow Lite через SavedModel или напрямую из модели Keras. См. примеры использования.

Из сохраненной модели

# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)

# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

Из модели Кераса

# build a Keras model
keras_model = build_keras_lstm(...)

# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

Пример

Keras LSTM для TensorFlow Lite Colab иллюстрирует комплексное использование интерпретатора TensorFlow Lite.

Поддерживаемые API TensorFlow RNN

Мы поддерживаем готовое преобразование Keras LSTM в TensorFlow Lite. Подробную информацию о том, как это работает, см. в интерфейсе Keras LSTM. и к логике преобразования здесь .

Также важно выделить контракт LSTM TensorFlow Lite относительно определения операции Keras:

  1. Размерность 0 входного тензора — это размер пакета.
  2. Размерность 0 тензора recurrent_weight — это количество выходов.
  3. Тензоры веса и recurrent_kernel транспонируются.
  4. Транспонированный вес, транспонированные recurrent_kernel и тензоры смещения разбиваются на 4 тензора одинакового размера по размерности 0. Они соответствуют входным воротам, воротам забывания, ячейке и выходным воротам .

Варианты Кераса LSTM

Время мажор

Пользователи могут выбрать временной режим или его отсутствие. Keras LSTM добавляет основной по времени атрибут в атрибуты def функции. Для однонаправленной последовательности LSTM мы можем просто сопоставить основной атрибут времени unidirecional_sequence_lstm.

Двунаправленный LSTM

Двунаправленный LSTM можно реализовать с помощью двух слоев Keras LSTM, одного для прямого и одного для обратного, примеры см. здесь . Как только мы видим атрибут go_backward, мы распознаем его как обратный LSTM, а затем группируем прямой и обратный LSTM вместе. Это будущая работа. В настоящее время это создает две операции UnidirectSequenceLSTM в модели TensorFlow Lite.

Примеры пользовательских преобразований LSTM

TensorFlow Lite также предоставляет способ преобразования пользовательских реализаций LSTM. Здесь мы используем LSTM от Lingvo в качестве примера того, как это можно реализовать. Подробности смотрите в интерфейсе lingvo.LSTMCellSimple и логике преобразования здесь . Мы также приводим пример другого определения LSTM Lingvo в интерфейсе lingvo.LayerNormalizedLSTMCellSimple и его логику преобразования здесь .

«Принесите свой собственный TensorFlow RNN» в TensorFlow Lite

Если интерфейс RNN пользователя отличается от поддерживаемых стандартом, есть несколько вариантов:

Вариант 1. Напишите код адаптера на Python TensorFlow, чтобы адаптировать интерфейс RNN к интерфейсу RNN Keras. Это означает tf.function с аннотацией tf_implements для сгенерированной функции интерфейса RNN, которая идентична той, которая сгенерирована слоем Keras LSTM. После этого будет работать тот же API преобразования, который использовался для Keras LSTM.

Вариант 2. Если вышеуказанное невозможно (например, в Keras LSTM отсутствуют некоторые функции, которые в настоящее время предоставляются с помощью объединенной операции LSTM TensorFlow Lite, подобной нормализации слоев), то расширьте конвертер TensorFlow Lite, написав собственный код преобразования, и подключите его к подготовке. -composite-functions MLIR-пройдите сюда . Интерфейс функции следует рассматривать как контракт API и должен содержать аргументы, необходимые для преобразования в объединенные операции TensorFlow Lite LSTM, т. е. ввод, смещение, веса, проекцию, нормализацию слоев и т. д. Предпочтительно, чтобы тензоры передавались в качестве аргументов этому функция должна иметь известный ранг (т. е. RankedTensorType в MLIR). Это значительно упрощает написание кода преобразования, который может принимать эти тензоры как RankedTensorType и помогает преобразовать их в ранжированные тензоры, соответствующие операндам объединенного оператора TensorFlow Lite.

Полным примером такого процесса преобразования является преобразование LSTMCellSimple в TensorFlow Lite от Lingvo.

LSTMCellSimple в Lingvo определен здесь . Модели, обученные с помощью этой ячейки LSTM, можно преобразовать в TensorFlow Lite следующим образом:

  1. Оберните все использование LSTMCellSimple в tf.function с аннотацией tf_implements, которая помечена как таковая (например, lingvo.LSTMCellSimple здесь будет хорошим именем аннотации). Убедитесь, что созданная tf.function соответствует интерфейсу функции, ожидаемой в коде преобразования. Это контракт между автором модели, добавляющим аннотацию, и кодом преобразования.

  2. Расширьте этап подготовки составных функций, чтобы подключить пользовательскую составную операцию к преобразованию операции Flowed LSTM TensorFlow Lite. См. код преобразования LSTMCellSimple .

    Конверсионный договор:

  3. Тензоры веса и проекции транспонируются.

  4. {Вход, рекуррентный} в {ячейку, входной вентиль, вентиль забывания, выходной вентиль} извлекаются путем разрезания транспонированного весового тензора.

  5. {Смещение} на {ячейку, входной вентиль, вентиль забывания, выходной вентиль} извлекается путем разрезания тензора смещения.

  6. Проекция извлекается путем разрезания транспонированного тензора проекции.

  7. Аналогичное преобразование написано для LayerNormalizedLSTMCellSimple .

  8. Оставшуюся часть инфраструктуры преобразования TensorFlow Lite, включая все определенные проходы MLIR , а также окончательный экспорт в плоский буфер TensorFlow Lite, можно использовать повторно.

Известные проблемы/ограничения

  1. В настоящее время поддерживается только преобразование Keras LSTM без сохранения состояния (поведение в Keras по умолчанию). Преобразование Keras LSTM с сохранением состояния — это будущая работа.
  2. По-прежнему возможно моделировать уровень Keras LSTM с сохранением состояния, используя базовый уровень Keras LSTM без сохранения состояния и явно управляя состоянием в пользовательской программе. Такую программу TensorFlow все равно можно преобразовать в TensorFlow Lite, используя описанную здесь функцию.
  3. Двунаправленный LSTM в настоящее время моделируется как две операции UnidirectSequenceLSTM в TensorFlow Lite. Это будет заменено одной операцией BidirectSequenceLSTM.