Visão geral
O TensorFlow Lite oferece suporte à conversão de modelos RNN do TensorFlow 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 a quantização.
Como existem muitas variantes de APIs RNN no TensorFlow, nossa abordagem foi dupla:
- Forneça suporte nativo para APIs RNN padrão do TensorFlow, como Keras LSTM. Esta é a opção recomendada.
- Forneça uma interface na infraestrutura de conversão para que implementações de RNN definidas pelo usuário sejam conectadas e convertidas para o TensorFlow Lite. Fornecemos alguns exemplos prontos para uso de tal conversão usando as interfaces LSTMCellSimple e LayerNormalizedLSTMCellSimple RNN do lingvo.
API do conversor
O recurso faz parte da versão TensorFlow 2.3. Também está disponível através do pip tf-nightly ou do head.
Essa funcionalidade de conversão está disponível ao converter para 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
Keras LSTM para TensorFlow Lite Colab ilustra o uso completo com o interpretador TensorFlow Lite.
APIs RNNs do TensorFlow compatíveis
Conversão Keras LSTM (recomendado)
Oferecemos suporte à conversão pronta para uso de Keras LSTM para 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 no que diz respeito à 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 de peso e recurrent_kernel são transpostos.
- Os tensores de peso transposto, recurrent_kernel transposto e polarização são divididos em 4 tensores de tamanhos iguais ao longo da dimensão 0. Eles correspondem à porta de entrada, porta de esquecimento, célula e porta de saída .
Variantes Keras LSTM
Tempo principal
Os usuários podem escolher o tempo maior ou nenhum tempo maior. Keras LSTM adiciona um atributo time-major nos atributos def da função. Para LSTM de sequência unidirecional, 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 . Depois de vermos o atributo go_backward, nós o reconhecemos como LSTM reverso e, em seguida, agrupamos LSTM para frente e para trás. Este é um trabalho futuro. Atualmente, isso cria duas operações UnidirecionalSequenceLSTM no modelo TensorFlow Lite.
Exemplos de conversão LSTM definidos pelo usuário
O TensorFlow Lite também fornece 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 definição 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 interfaces padrão suportadas, há algumas opções:
Opção 1: Escreva o código do adaptador em TensorFlow python para adaptar a interface RNN à interface Keras RNN. Isso significa uma anotação tf.function com 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 normalização de camada semelhante à operação LSTM fundida do TensorFlow Lite), estenda o conversor do TensorFlow Lite escrevendo um código de conversão personalizado e conecte-o ao prepare -composite-functions MLIR-passe 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. função para ter classificação conhecida (ou seja, RankedTensorType em MLIR). Isso torna muito mais fácil escrever código de conversão que pode assumir esses tensores como RankedTensorType e ajuda a transformá-los em tensores classificados correspondentes aos operandos fundidos do operador TensorFlow Lite.
Um exemplo completo desse fluxo de conversão é a conversão LSTMCellSimple para TensorFlow Lite do Lingvo.
O LSTMCellSimple no Lingvo é definido aqui . Os modelos treinados com esta célula LSTM podem ser convertidos para 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 tf.function gerado corresponda à 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:
Tensores de peso e projeção são transpostos.
A {entrada, recorrente} para {célula, porta de entrada, porta de esquecimento, porta de saída} são extraídas cortando o tensor de peso transposto.
A {polarização} para {célula, porta de entrada, porta de esquecimento, porta de saída} é extraída cortando o tensor de polarização.
A projeção é extraída cortando o tensor de projeção transposto.
Uma conversão semelhante foi escrita para LayerNormalizedLSTMCellSimple .
O restante da infraestrutura de conversão do TensorFlow Lite, incluindo todas as passagens MLIR definidas, bem como a exportação final para o flatbuffer do TensorFlow Lite, pode ser reutilizada.
Problemas/limitações conhecidos
- Atualmente há suporte apenas para conversão de Keras LSTM sem estado (comportamento padrão em 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 para TensorFlow Lite usando o recurso descrito aqui.
- O LSTM bidirecional é atualmente modelado como duas operações UnidirecionalSequenceLSTM no TensorFlow Lite. Isso será substituído por uma única operação BidirecionalSequenceLSTM.