Conversão da RNN do TensorFlow para o TensorFlow Lite

Visão geral

O TensorFlow Lite oferece suporte à conversão de modelos de RNN do TensorFlow para as operações de LSTM fundidas do TensorFlow Lite. As operações combinadas existem para maximizar o desempenho das implementações do kernel, além de fornecer uma interface de nível superior para definir transformações complexas, como a quantização.

Como existem muitas variantes das APIs de RNN no TensorFlow, nossa abordagem foi duplicada:

  1. Fornecer suporte nativo para APIs RNN padrão do TensorFlow, como Keras LSTM. Essa é a opção recomendada.
  2. Fornecer uma interface na infraestrutura de conversão para implementações de RNN definidas pelo usuário para conexão e conversão para o TensorFlow Lite. Fornecemos alguns exemplos prontos dessa conversão usando as interfaces NN LSTMCellSimple e LayerNormalizedLSTMCellSimple do lingvo.

API Converter

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

Essa funcionalidade de conversão está disponível ao fazer a conversão para o TensorFlow Lite usando 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

O Keras LSTM para o TensorFlow Lite Colab ilustra o uso completo com o intérprete do TensorFlow Lite.

APIs de RNNs do TensorFlow com suporte

Oferecemos suporte à conversão imediata do Keras LSTM para o TensorFlow Lite. Para detalhes sobre como isso funciona, consulte a interface Keras LSTM e a lógica de conversão aqui.

Também é importante destacar o contrato LSTM do TensorFlow Lite em relação à definição da operação do Keras:

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

Variantes do Keras LSTM

Horário principal

Os usuários podem escolher o período maior ou nenhum. A Keras LSTM adiciona um atributo de tempo maior nos atributos def def. Para a LSTM de sequência unidirecional, podemos simplesmente mapear para o atributo time principal de unidirecional_sequence_lstm.

LSTM bidirecional

O LSTM bidirecional pode ser implementado com duas camadas do Keras LSTM, uma para avançada e outra para trás. Confira exemplos aqui. Quando vemos o atributo go_backward, o reconhecemos como LSTM para trás e, em seguida, agrupamos os LSTMs para frente e para trás. Este é um trabalho futuro. Atualmente, isso cria duas operações UnidirectionalSequenceLSTM no modelo do TensorFlow Lite.

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

O TensorFlow Lite também oferece uma maneira de converter implementações de LSTM definidas pelo usuário. Aqui, usamos o LSTM do Lingvo como 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 de outra definição de LSTM do Lingvo na interface lingvo.LayerNormalizedLSTMCellSimple (em inglês) e a lógica de conversão aqui (links em inglês).

"Traga sua própria RNN do TensorFlow" para o TensorFlow Lite

Se a interface RNN de um usuário for diferente das compatíveis padrão, há algumas opções:

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

Opção 2:se não for possível fazer isso acima, por exemplo, faltam algumas funcionalidades que estão atualmente expostas pela operação LSTM fundida do TensorFlow Lite, como a normalização de camadas. Em seguida, estenda o conversor do TensorFlow Lite escrevendo código de conversão personalizado e conecte-o ao aqui (em inglês) do prepare-composite-functions. A interface da função precisa ser tratada como um contrato de API e conter os argumentos necessários para converter em operações LSTM do TensorFlow Lite fundidas, ou seja, entrada, polarização, pesos, projeção, normalização de camada etc. É preferível que os tensores transmitidos como argumentos para essa função tenham classificação conhecida (ou seja, RankedTensorType em MLIR). Isso facilita muito a gravação de um código de conversão que possa presumir esses tensores como RankedTensorType e ajuda a transformá-los em tensores classificados correspondentes aos operandos do operador do TensorFlow Lite combinados.

Um exemplo completo desse fluxo de conversão é a conversão de LSTMCellSimple para TensorFlow Lite da Lingvo.

O LSTMCellSimple no Lingvo é definido aqui. Os modelos treinados com essa célula LSTM podem ser convertidos para o TensorFlow Lite da seguinte maneira:

  1. Encapsule todos os usos de LSTMCellSimple em uma tf.function com uma anotação tf_implements rotulada dessa forma (por exemplo, lingvo.LSTMCellSimple seria um bom nome de anotação aqui). Verifique se a função tf.function gerada corresponde à interface da função esperada no código de conversão. Esse é um contrato entre o autor do modelo que adiciona a anotação e o código de conversão.
  2. Estenda a passagem das funções de preparação compostas para conectar uma operação composta personalizada à conversão de operações LSTM combinadas do TensorFlow Lite. Consulte o código de conversão LSTMCellSimple.

    Contrato de conversão:

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

  4. O {input, recurrent} para {cell, input gate, remember gate, output gate} são extraídos ao dividir o tensor de peso transposto.

  5. O {bias} para {cell, input gate, remember gate, output gate} são extraídos ao dividir o tensor de viés.

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

  7. Conversão semelhante é gravada para LayerNormalizedLSTMCellSimple.

  8. É possível reutilizar o restante da infraestrutura de conversão do TensorFlow Lite, incluindo todas as passagens MLIR definidas e a exportação final para o flatbuffer do TensorFlow Lite.

Problemas/limitações conhecidos

  1. Atualmente, há suporte apenas para a conversão de Keras LSTM sem estado (comportamento padrão no Keras). A conversão do Keras LSTM com estado é um trabalho futuro.
  2. Ainda é possível modelar uma camada Keras LSTM com estado usando a camada subjacente do Keras LSTM sem estado e gerenciando o estado explicitamente no programa do usuário. Esse programa do TensorFlow ainda pode ser convertido para o TensorFlow Lite usando o recurso descrito aqui.
  3. Atualmente, o LSTM bidirecional é modelado como duas operações UnidirectionalSequenceLSTM no TensorFlow Lite. Ela será substituída por uma única op BidirectionalSequenceLSTM.