Visão geral
O LiteRT é compatível com a conversão de modelos RNN do TensorFlow em operações LSTM fundidas do LiteRT. As operações de fusão existem para maximizar a performance das implementações de kernel subjacentes e fornecer uma interface de nível mais alto para definir transformações complexas, como quantização.
Como há muitas variantes de APIs de RNN no TensorFlow, nossa abordagem foi dupla:
- Oferecer suporte nativo para APIs RNN padrão do TensorFlow, como LSTM do Keras. Essa é a opção recomendada.
- Forneça uma interface para a infraestrutura de conversão de implementações de RNN definidas pelo usuário para conectar e converter em LiteRT. Oferecemos alguns exemplos prontos de conversão usando as interfaces LSTMCellSimple e LayerNormalizedLSTMCellSimple do lingvo.
API Converter
O recurso faz parte da versão 2.3 do TensorFlow. Ele também está disponível no pip tf-nightly ou no head.
Essa funcionalidade de conversão está disponível ao converter para LiteRT usando um SavedModel ou diretamente do modelo Keras. Confira exemplos de uso.
De um 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()
De um 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
LSTM do Keras para LiteRT O Colab ilustra o uso de ponta a ponta com o intérprete LiteRT.
APIs de RNNs do TensorFlow compatíveis
Conversão de LSTM do Keras (recomendada)
Oferecemos suporte à conversão imediata de LSTM do Keras para LiteRT. Para detalhes sobre como isso funciona, consulte a interface LSTM do Keras e a lógica de conversão aqui.
Também é importante destacar o contrato LSTM do LiteRT em relação à definição da operação do Keras:
- A dimensão 0 do tensor input é 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.
- Os tensores de peso transposto, recurrent_kernel transposto e bias são divididos em quatro tensores de tamanho igual ao longo da dimensão 0. Elas correspondem a entrada, esquecimento, célula e saída.
Variantes de LSTM do Keras
Tempo principal
Os usuários podem escolher a dimensão de tempo ou nenhuma. A LSTM do Keras adiciona um atributo de tempo principal nos atributos de definição de função. Para LSTM de sequência unidirecional, podemos simplesmente mapear para o atributo principal de tempo de unidirecional_sequence_lstm.
LSTM bidirecional
A LSTM bidirecional pode ser implementada com duas camadas LSTM do Keras, uma para frente e outra para trás. Confira exemplos aqui. Quando vemos o atributo "go_backward", reconhecemos como LSTM invertida e agrupamos LSTM direta e invertida. Isso é trabalho futuro. No momento, isso cria duas operações UnidirectionalSequenceLSTM no modelo LiteRT.
Exemplos de conversão de LSTM definidos pelo usuário
O LiteRT também oferece uma maneira de converter implementações de LSTM definidas pelo usuário. Aqui usamos a 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 e a lógica de conversão aqui.
"Traga sua própria RNN do TensorFlow" para o LiteRT
Se a interface RNN de um usuário for diferente das padrão compatíveis, há algumas opções:
Opção 1:escreva um código de adaptador em Python do TensorFlow para adaptar a interface da RNN à interface da RNN do Keras. Isso significa uma tf.function com anotação tf_implements 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 LSTM do Keras vai funcionar.
Opção 2:se o acima não for possível (por exemplo, a LSTM do Keras não tiver alguma funcionalidade exposta atualmente pela operação de LSTM fundida do LiteRT, como normalização de camada), estenda o conversor do LiteRT escrevendo um código de conversão personalizado e conecte-o à transmissão MLIR prepare-composite-functions aqui. A interface da função precisa ser tratada como um contrato de API e conter os argumentos necessários para a conversão em operações LSTM LiteRT fundidas, ou seja, entrada, viés, 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 código de conversão que pode assumir esses tensores como RankedTensorType e ajuda a transformá-los em tensores classificados correspondentes aos operandos do operador LiteRT fundido.
Um exemplo completo desse fluxo de conversão é a conversão de LSTMCellSimple do Lingvo para LiteRT.
A LSTMCellSimple no Lingvo é definida aqui. Os modelos treinados com essa célula LSTM podem ser convertidos para LiteRT da seguinte maneira:
- Encapsule todos os usos de LSTMCellSimple em uma tf.function com uma anotação tf_implements rotulada como tal (por exemplo, lingvo.LSTMCellSimple seria um bom nome de anotação aqui). Verifique se a tf.function gerada 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 a transmissão prepare-composite-functions para conectar uma operação de fusão LSTM personalizada do LiteRT. Consulte o código de conversão LSTMCellSimple.
O contrato de conversão:
Os tensores de peso e projeção são transpostos.
Os {input, recurrent} para {cell, input gate, forget gate, output gate} são extraídos dividindo o tensor de peso transposto.
O {bias} para {cell, input gate, forget gate, output gate} é extraído ao segmentar o tensor de bias.
A projeção é extraída dividindo o tensor de projeção transposto.
Uma conversão semelhante é gravada para LayerNormalizedLSTMCellSimple.
O restante da infraestrutura de conversão do LiteRT, incluindo todas as transmissões do MLIR definidas, bem como a exportação final para o flatbuffer do LiteRT, pode ser reutilizado.
Limitações e problemas conhecidos
- No momento, só é possível converter LSTM do Keras sem estado (comportamento padrão no Keras). A conversão de LSTM do Keras com estado é um trabalho futuro.
- Ainda é possível modelar uma camada LSTM do Keras com estado usando a camada LSTM do Keras sem estado subjacente e gerenciando o estado explicitamente no programa do usuário. Um programa do TensorFlow ainda pode ser convertido para LiteRT usando o recurso descrito aqui.
- No momento, o LSTM bidirecional é modelado como duas operações UnidirectionalSequenceLSTM no LiteRT. Isso será substituído por uma única operação BidirectionalSequenceLSTM.