Преобразование 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. Она также доступна через pip-пакет tf-nightly или через 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()

Из модели Keras 

# 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 to 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. Они соответствуют входному вентилю, вентилю забывания, ячейке и выходному вентилю .

Варианты Keras LSTM

Время мажор

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

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

Двунаправленную LSTM-сеть можно реализовать с помощью двух LSTM-слоёв Keras: одного для прямого и одного для обратного. Примеры см. здесь . Как только мы видим атрибут go_backward, мы распознаём её как обратную LSTM-сеть, а затем группируем прямую и обратную LSTM-сеть вместе. Это — задача на будущее. В настоящее время это создаёт две операции UnidirectionSequenceLSTM в модели LiteRT.

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

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

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

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

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

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

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

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

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

  2. Расширьте проход prepare-composite-functions для подключения пользовательского композитного оператора к преобразованию LSTM-оператора LiteRT. См. код преобразования 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 моделируется двумя операциями UnidirectionSequenceLSTM. В дальнейшем она будет заменена одной операцией BidirectionSequenceLSTM.