Visão geral
O TensorFlow Lite oferece suporte à conversão de modelos do TensorFlow RNN em operações LSTM fundidas do TensorFlow Lite. As operações fundidas existem para maximizar o desempenho de suas implementações de kernel subjacentes, bem como fornecer uma interface de nível superior para definir transformações complexas como quantização.
Como existem muitas variantes de APIs RNN no TensorFlow, nossa abordagem foi dupla:
- Forneça suporte nativo para APIs padrão do TensorFlow RNN, como Keras LSTM. Esta é a opção recomendada.
- Forneça uma interface na infraestrutura de conversão para implementações de RNN definidas pelo usuário para serem conectadas e convertidas para o TensorFlow Lite. Fornecemos alguns exemplos prontos de tal conversão usando as interfaces LSTMCellSimple e LayerNormalizedLSTMellSimple RNN do lingvo.
API do conversor
O recurso faz parte da versão 2.3 do TensorFlow. Também está disponível através do tf-nightly pip ou do head.
Essa funcionalidade de conversão está disponível ao converter para o TensorFlow Lite por meio de um SavedModel ou diretamente do modelo Keras. Veja 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 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 to TensorFlow Lite Colab ilustra o uso de ponta a ponta com o interpretador TensorFlow Lite.
APIs de RNNs do TensorFlow compatíveis
Conversão Keras LSTM (recomendado)
Oferecemos suporte à conversão pronta para uso do Keras LSTM para o TensorFlow Lite. Para obter detalhes sobre como isso funciona, consulte a interface Keras LSTM e para a lógica de conversão aqui .
Também é importante destacar o contrato LSTM do TensorFlow Lite com relação à definição da operação Keras:
- A dimensão 0 do tensor de entrada é o tamanho do lote.
- A dimensão 0 do tensor recurrent_weight é o número de saídas.
- Os tensores weight e recurrent_kernel são transpostos.
- O peso transposto, recurrent_kernel transposto e tensores de polarização são divididos em 4 tensores de tamanhos iguais ao longo da dimensão 0. Eles correspondem ao portão de entrada, portão de esquecimento, célula e portão de saída .
Variantes Keras LSTM
Tempo maior
Os usuários podem escolher tempo maior ou nenhum tempo maior. Keras LSTM adiciona um atributo de tempo maior nos atributos def da função. Para a sequência unidirecional LSTM, podemos simplesmente mapear para o atributo principal de tempo de unidirecional_sequence_lstm .
LSTM Bidirecional
O LSTM bidirecional pode ser implementado com duas camadas Keras LSTM, uma para frente e outra para trás, veja exemplos aqui . Uma vez que vemos o atributo go_backward, nós o reconhecemos como LSTM para trás, então agrupamos o LSTM para frente e para trás. Este é um trabalho futuro. Atualmente, isso cria duas operações UnidirectionalSequenceLSTM no modelo TensorFlow Lite.
Exemplos de conversão LSTM definidos pelo usuário
O TensorFlow Lite também oferece uma maneira de converter implementações LSTM definidas pelo usuário. Aqui usamos o LSTM do Lingvo como um exemplo de como isso pode ser implementado. Para obter detalhes, consulte a interface lingvo.LSTMCellSimple e a lógica de conversão aqui . Também fornecemos um exemplo para outra das definições LSTM do Lingvo na interface lingvo.LayerNormalizedLSTMCellSimple e sua lógica de conversão aqui .
“Traga seu próprio TensorFlow RNN” para o TensorFlow Lite
Se a interface RNN de um usuário for diferente das suportadas padrão, há algumas opções:
Opção 1: escreva o código do adaptador em python do TensorFlow para adaptar a interface RNN à interface Keras RNN. Isso significa uma tf.function com anotação tf_implements na função da interface RNN gerada que é idêntica àquela gerada pela camada Keras LSTM. Depois disso, a mesma API de conversão usada para Keras LSTM funcionará.
Opção 2: Se o acima não for possível (por exemplo, o Keras LSTM está faltando alguma funcionalidade que está atualmente exposta pela operação LSTM fundida do TensorFlow Lite como normalização de camada), então estenda o conversor do TensorFlow Lite escrevendo o código de conversão personalizado e conecte-o ao prepare -composite-functions MLIR-pass aqui . A interface da função deve ser tratada como um contrato de API e deve conter os argumentos necessários para converter em operações fundidas do TensorFlow Lite LSTM - ou seja, entrada, polarização, pesos, projeção, normalização de camada, etc. É preferível para os tensores passados como argumentos para isso função ter uma classificação conhecida (ou seja, RankTensorType em MLIR). Isso torna muito mais fácil escrever o código de conversão que pode assumir esses tensores como RankTensorType e ajuda a transformá-los em tensores classificados correspondentes aos operandos do operador do TensorFlow Lite fundido.
Um exemplo completo desse fluxo de conversão é a conversão LSTMCellSimple para TensorFlow Lite da Lingvo.
O LSTMCellSimple no Lingvo é definido aqui . Os modelos treinados com esta célula LSTM podem ser convertidos para o TensorFlow Lite da seguinte forma:
- Envolva todos os usos de LSTMCellSimple em um tf.function com uma anotação tf_implements rotulada como tal (por exemplo, lingvo.LSTMCellSimple seria um bom nome de anotação aqui). Certifique-se de que o tf.function gerado corresponde à interface da função esperada no código de conversão. Este é um contrato entre o autor do modelo que adiciona a anotação e o código de conversão.
Estenda o passo prepare-composite-functions para conectar uma operação composta personalizada à conversão de operação LSTM fundida do TensorFlow Lite. Consulte o código de conversão LSTMCellSimple .
O contrato de conversão:
Os tensores de peso e projeção são transpostos.
O {input, recurrent} para {cell, input gate, forget gate, output gate} são extraídos fatiando o tensor de peso transposto.
O {bias} para {cell, input gate, forget gate, output gate} são extraídos fatiando o tensor de polarização.
A projeção é extraída fatiando o tensor de projeção transposta.
Uma conversão semelhante é gravada para LayerNormalizedLSTMCellSimple .
O restante da infraestrutura de conversão do TensorFlow Lite, incluindo todos os passes MLIR definidos, bem como a exportação final para o flatbuffer do TensorFlow Lite, pode ser reutilizado.
Problemas/limitações conhecidos
- Atualmente, há suporte apenas para converter Keras LSTM sem estado (comportamento padrão no Keras). A conversão Stateful Keras LSTM é um trabalho futuro.
- Ainda é possível modelar uma camada Keras LSTM com estado usando a camada Keras LSTM sem estado subjacente e gerenciando o estado explicitamente no programa do usuário. Esse programa TensorFlow ainda pode ser convertido em TensorFlow Lite usando o recurso descrito aqui.
- Atualmente, o LSTM bidirecional é modelado como duas operações UnidirectionalSequenceLSTM no TensorFlow Lite. Isso será substituído por uma única operação BidirectionalSequenceLSTM.