Conversão de RNN do TensorFlow para LiteRT

Visão geral

A LiteRT oferece suporte à conversão de modelos de RNN do TensorFlow em operações LSTM combinadas. As operações fundidas existem para maximizar o desempenho as implementações do kernel, além de oferecer um nível mais alto para definir transformações complexas, como a quantização.

Como há muitas variantes de APIs de RNN no TensorFlow, nossa abordagem foi duas vezes:

  1. Ofereça suporte nativo para APIs RNN padrão do TensorFlow, como o Keras LSTM. Essa é a opção recomendada.
  2. Forneça uma interface para a infraestrutura de conversão de Implementações de RN definidas pelo usuário para conexão e conversão LiteRT. Fornecemos alguns exemplos imediatos desses conversão usando linguagem LSTMCellSimple e LayerNormalizedLSTMCellSimple RNN.

API Converter

O recurso faz parte da versão 2.3 do TensorFlow. Ele também está disponível na do pip tf-nightly ou do cabeçalho.

Essa funcionalidade de conversão está disponível ao converter para LiteRT por um SavedModel ou diretamente do modelo do Keras. Confira exemplos de uso.

Do modelo salvo

# 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()

Do modelo do Keras

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

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

Exemplo

Keras LSTM para LiteRT Colab (em inglês) ilustra o uso completo com o intérprete LiteRT.

APIs de RNNs do TensorFlow com suporte

Oferecemos suporte à conversão pronta para uso do Keras LSTM para LiteRT. Para detalhes sobre como isso funciona, consulte o Interface Keras LSTM e à lógica de conversão aqui.

Também é importante destacar o contrato de LSTM da LiteRT com respeito à definição da operação Keras:

  1. A dimensão 0 do tensor de input é o tamanho do lote.
  2. A dimensão 0 do tensor recurrent_weight é o número de de entrada e saída.
  3. Os tensores weight e recurrent_kernel são transpostos.
  4. Os tensores de peso transposto, recurrent_kernel transposto e bias são divididos em 4 tensores de tamanho igual ao longo da dimensão 0. Estes correspondem a portas de entrada, de esquecimento, de célula e de saída.

Variantes do Keras LSTM

Maior tempo

Os usuários podem escolher tempo maior ou nenhum tempo maior. A Keras LSTM nos atributos def de função. Para a sequência unidirecional LSTM, pode simplesmente ser mapeada para o atributo principal de tempo.

LSTM bidirecional

O LSTM bidirecional pode ser implementado com duas camadas do Keras LSTM, uma para para frente e outro para trás. Veja exemplos aqui. Ao vermos o atributo go_backward, nós o reconhecemos como LSTM para trás, então agrupamos adiante e LSTM para trás juntas. Esse é um trabalho futuro. Atualmente, isso cria duas operações UnidirectionalSequenceLSTM na interface um modelo de machine learning.

Exemplos de conversão de LSTM definidos pelo usuário

A LiteRT também oferece uma maneira de converter o LSTM definido pelo usuário. e implementações. Aqui, usamos o LSTM de Lingvo como um exemplo de como isso pode ser implementado. Para mais detalhes, consulte a Interface lingvo.LSTMCellSimple e a lógica de conversão aqui. Também fornecemos um exemplo para outra definição de LSTM da Lingvo em Interface lingvo.LayerNormalizedLSTMCellSimple e a lógica de conversão aqui.

"traga sua própria RNN do TensorFlow" para o LiteRT

Se a interface RNN do usuário for diferente das padrão com suporte, haverá há algumas opções:

Opção 1:escrever o código do adaptador no TensorFlow Python para adaptar a interface RNN para a interface RNN do Keras. Isso significa que uma tf.function com tf_implements anotação em da função da interface RNN gerada é idêntica à gerada pelo a camada Keras LSTM. Depois disso, a mesma API de conversão usada para o Keras LSTM vai funcionar.

Opção 2:se a opção acima não for possível (por exemplo, o LSTM do Keras está sem alguns funcionalidade atualmente exposta pela operação LSTM combinada do LiteRT, como normalização de camadas) e depois estenda o conversor LiteRT escrevendo código de conversão personalizado e conectá-lo às funções prepare-composite MLIR-pass aqui. A interface da função deve ser tratada como um contrato de API e deve contêm os argumentos necessários para converter para a versão fundida do LiteRT LSTM. operações, como entrada, viés, pesos, projeção, normalização de camadas etc. é preferível que os tensores passados como argumentos para essa função tenham conhecimento classificação (ou seja, RankedTensorType em MLIR). Isso torna muito mais fácil escrever que pode assumir esses tensores como RankedTensorType e ajuda transformá-los em tensores classificados correspondentes à função fundida do LiteRT operandos do operador.

Um exemplo completo desse fluxo de conversão é o LSTMCellSimple da Lingvo Conversão LiteRT.

O LSTMCellSimple no Lingvo é definido aqui. Os modelos treinados com essa célula LSTM podem ser convertidos para LiteRT como da seguinte forma:

  1. Una todos os usos de LSTMCellSimple em uma tf.function com um rotulada dessa forma (por exemplo, lingvo.LSTMCellSimple seria uma bom nome de anotação aqui). Confira se a tf.function que foi gerada corresponde à interface da função esperada no código de conversão. Isso é um contrato entre o autor do modelo, que adiciona a anotação e o o código de conversão.
  2. Estender a passagem de prepare-composite-functions para conectar uma operação composta personalizada para a conversão de op LSTM combinada com LiteRT. Consulte LSTMCellSimple o código de conversão.

    O contrato de conversão:

  3. Os tensores de peso e projeção são transpostos.

  4. A entrada {input, recurrent} para {cell, input gate, remember gate, output} gate} são extraídos fracionando o tensor de peso transposto.

  5. Os valores de {bias} a {bias} são é extraído ao dividir o tensor de viés.

  6. A projeção é extraída ao dividir o tensor de projeção transposto.

  7. A conversão semelhante é escrita para LayerNormalizedLSTMCellSimple.

  8. O restante da infraestrutura de conversão da LiteRT, incluindo todas as Passes MLIR definido, assim como a exportação final para flatbuffer LiteRT pode ser reutilizadas.

Limitações/problemas conhecidos

  1. No momento, só é possível converter o Keras LSTM sem estado (padrão no Keras). A conversão LSTM do Keras com estado é um trabalho futuro.
  2. Ainda é possível modelar uma camada Keras LSTM com estado usando o camada Keras LSTM sem estado subjacente e gerenciando o estado explicitamente programa do usuário. Esse programa do TensorFlow ainda pode ser convertido em LiteRT usando o recurso descrito aqui.
  3. O LSTM bidirecional atualmente é modelado como dois UnidirectionalSequenceLSTM operações na LiteRT. Ela será substituída por uma única op. BidirectionalSequenceLSTM