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

Обзор

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

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

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

API конвертера

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

Эта функция преобразования доступна при преобразовании в LiteRT через 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 для LiteRT Colab иллюстрирует комплексное использование интерпретатора LiteRT.

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

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

Также важно выделить контракт LSTM LiteRT в отношении определения операции 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 в модели LiteRT.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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