Visão geral
Esta página descreve o projeto e as etapas necessárias para converter operações compostas no TensorFlow para mesclar operações no LiteRT. Essa infraestrutura é de uso geral e aceita a conversão de qualquer operação composta no TensorFlow para uma operação combinada correspondente no LiteRT.
Um exemplo de uso dessa infraestrutura é a fusão de operações de RNN do TensorFlow para LiteRT, conforme detalhado aqui.
O que são operações fundidas
As operações do TensorFlow podem ser primitivas, por exemplo, tf.add ou podem ser compostas de outras operações primitivas, por exemplo, tf.einsum (em inglês). Um primitivo aparece como um único nó no gráfico do TensorFlow, enquanto um é uma coleção de nós no gráfico do TensorFlow. A execução de um uma operação composta equivale à execução de cada um dos primitivos as operações.
Uma operação fundida corresponde a uma única operação que engloba todas as computação realizada por cada operação primitiva dentro da operação composta.
Benefícios das operações fundidas
Existem operações fundidas para maximizar o desempenho do kernel subjacente de código aberto, otimizando a computação geral e reduzindo pegada de carbono. Isso é muito valioso, especialmente para cargas de trabalho de inferência de baixa latência e plataformas para dispositivos móveis com recursos limitados.
As operações fundidas também oferecem uma interface de nível superior para definir transformações como a quantização, o que seria inviável ou muito em um nível mais granular.
A LiteRT tem muitas instâncias de operações fundidas pelos motivos articulada acima. Essas operações fundidas normalmente correspondem a operações no programa TensorFlow de origem. Exemplos de operações compostas em TensorFlow que são implementados como uma única operação combinada no LiteRT incluem várias operações de RNN, como sequência unidirecional e bidirecional LSTM, convolução (conv2d, bias add, relu), totalmente conectada (matmul, bias add, Relu) e muito mais. No LiteRT, a quantização LSTM só está disponível no momento implementadas nas operações de LSTM fundidas.
Desafios com operações fundidas
Como converter operações compostas do TensorFlow em operações fundidas no A LiteRT é um problema difícil. Isso ocorre pelos seguintes motivos:
As operações compostas são representadas no gráfico do TensorFlow como um conjunto de operações primitivas sem um limite bem definido. Pode ser muito difícil de identificar (por exemplo, por meio de correspondência de padrões) do subgráfico que correspondem a uma operação composta.
Pode haver mais de uma implementação do TensorFlow direcionada a um dispositivo Operação LiteRT. Por exemplo, há muitas implementações de LSTM no TensorFlow (Keras, Babelfish/lingvo etc.), e cada um deles é composto por operações primitivas diferentes, mas todas elas ainda poderiam ser convertidas mesma operação LSTM combinada no LiteRT.
Assim, a conversão de operações combinadas tem se mostrado bastante desafiador.
Como converter de uma operação composta para uma operação personalizada do TFLite (recomendado)
Unir a operação composta em um tf.function
Em muitos casos, parte do modelo pode ser mapeada para uma única operação em
TFLite. Isso ajuda na performance ao escrever uma implementação otimizada
para operações específicas. Para criar uma operação fundida no TFLite,
identificar a parte do gráfico que representa uma operação fundida e envolvê-la
tf.function com
"experimental_implements" a um tf.function, que tem atributos
valor tfl_fusable_op com valor true. Se a operação personalizada demorar
e passá-los como parte da mesma "experimental_implements".
Por exemplo:
def get_implements_signature():
implements_signature = [
# 'name' will be used as a name for the operation.
'name: "my_custom_fused_op"',
# attr "tfl_fusable_op" is required to be set with true value.
'attr {key: "tfl_fusable_op" value { b: true } }',
# Example attribute "example_option" that the op accepts.
'attr {key: "example_option" value { i: %d } }' % 10
]
return ' '.join(implements_signature)
@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
# An empty function that represents pre/post processing example that
# is not represented as part of the Tensorflow graph.
output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
return output_1, output_2
class TestModel(tf.Module):
def __init__(self):
super(TestModel, self).__init__()
self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
@tf.function(input_signature=[
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
])
def simple_eval(self, input_a, input_b):
return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))
Não é necessário definir allow_custom_ops no conversor como
tfl_fusable_op já indicam isso.
Implementar uma operação personalizada e fazer o registro com o TFLite Interpreter
Implemente a operação fundida como uma operação personalizada do TFLite. Consulte instruções.
O nome de registro da operação precisa ser semelhante ao nome
especificado no atributo name na assinatura de implementação.
Um exemplo para a operação no exemplo é
TfLiteRegistration reg = {};
// This name must match the name specified in the implements signature.
static constexpr char kOpName[] = "my_custom_fused_op";
reg.custom_name = kOpName;
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, ®);
Como converter de operação composta para fusão (avançado)
A arquitetura geral para converter operações compostas do TensorFlow em As operações combinadas de LiteRT estão abaixo:
Unir a operação composta em um tf.function
No código-fonte do modelo do TensorFlow, identifique e abstraia o composto em um tf.function com o experimental_implements anotação de função. Confira um exemplo de pesquisa de incorporação. A define a interface e seus argumentos devem ser usados para implementar a lógica de conversão.
Escrever código de conversão
O código de conversão é escrito de acordo com a interface da função com o
implements. Veja um exemplo de fusão para embedding
pesquisa. Conceitualmente, o código de conversão substitui
implementação dessa interface com a função fundida.
Na passagem prepare-composite-functions, use o plug-in em seu objeto conversion ou código-fonte.
Em usos mais avançados, é possível implementar transformações complexas de os operandos da operação composta para derivar os operandos da função operação Consulte Keras LSTM. como exemplo.
Converter para LiteRT
Use o TFLiteConverter.from_saved_model API para converter em LiteRT.
Configurações avançadas
Agora descrevemos os detalhes de alto nível do design geral na conversão para fusão operações na LiteRT.
Como criar operações no TensorFlow
O uso de tf.function com o experimental_implements atributo de função permite que os usuários componham novas operações explicitamente usando operações primitivas do TensorFlow e especificar a interface uma operação composta é implementada. Isso é muito útil, pois fornece:
- Um limite bem definido para a operação composta na infraestrutura Gráfico do TensorFlow.
- Especifique explicitamente a interface implementada pela operação. A da função tf.function correspondem aos argumentos dessa interface.
Como exemplo, vamos considerar uma operação composta definida para implementar pesquisa de incorporação. Esse valor é mapeado para uma operação fundida no LiteRT.
@tf.function(
experimental_implements="embedding_lookup")
def EmbFprop(embs, ids_vec):
"""Embedding forward prop.
Effectively, it computes:
num = size of ids_vec
rets = zeros([num, embedding dim])
for i in range(num):
rets[i, :] = embs[ids_vec[i], :]
return rets
Args:
embs: The embedding matrix.
ids_vec: A vector of int32 embedding ids.
Returns:
The result of embedding lookups. A matrix of shape
[num ids in ids_vec, embedding dims].
"""
num = tf.shape(ids_vec)[0]
rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))
def EmbFpropLoop(i, embs, ids_vec, rets):
# row_id = ids_vec[i]
row_id = tf.gather(ids_vec, i)
# row = embs[row_id]
row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
# rets[i] = row
rets = inplace_ops.alias_inplace_update(rets, [i], row)
return embs, ids_vec, rets
_, _, rets = functional_ops.For(
start=0,
limit=num,
delta=1,
inputs=[embs, ids_vec, rets],
body=EmbFpropLoop,
rewrite_with_while=compiled)
if len(weight_shape) > 2:
rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
return rets
Fazendo com que os modelos usem operações compostas via tf.function como ilustrado acima, é possível construir uma infraestrutura geral para identifique e converta essas operações em operações LiteRT fundidas.
Extensão do conversor de LiteRT
O conversor LiteRT lançado no início deste ano só oferecia suporte a importar modelos do TensorFlow como um gráfico com todas as variáveis substituídas pelos respectivos valores constantes correspondentes. Isso não funciona para fusão de operações, como esses gráficos têm todas as funções embutidas, as variáveis podem ser transformadas constantes.
Para aproveitar
tf.function com o
experimental_implements durante o processo de conversão, as funções
precisam ser preservados até um
momento posterior do processo de conversão.
Por isso, implementamos um novo fluxo de trabalho de importação e conversão do TensorFlow modelos no conversor para oferecer suporte ao caso de uso de fusão de operações compostas. Mais especificamente, os novos recursos adicionados são:
- Importação de modelos salvos do TensorFlow para MLIR (em inglês)
- fundir operações compostas
- análise de mutabilidade variável
- congelar todas as variáveis somente leitura
Com isso, é possível realizar a fusão de operações usando as funções que representam as operações compostas antes da função inline e do congelamento variável.
Como implementar a fusão de operações
Vamos analisar a passagem de fusão da operação com mais detalhes. Esse cartão faz seguintes:
- Faça um loop por todas as funções no módulo MLIR.
- Se uma função tiver o atributo tf._implements, com base no atributo chama o utilitário de fusão da operação adequada.
- O utilitário de fusão da operação opera nos operandos e (que servem como a interface para a conversão) e substitui o corpo da função com um corpo de função equivalente contendo a operação fundida.
- Em muitos casos, o corpo substituído conterá operações diferentes da operação fundida. Elas correspondem a algumas transformações estáticas na operandos da função para obter os operandos da operação fundida. Como todos esses cálculos podem ser dobrados com frequência, eles não são presente no flatbuffer exportado, em que apenas a operação fundida existem.
Este é o snippet de código do cartão que mostra o fluxo de trabalho principal:
void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
StringAttr attr) {
if (attr.getValue() == "embedding_lookup") {
func.eraseBody();
func.addEntryBlock();
// Convert the composite embedding_lookup function body to a
// TFLite fused embedding_lookup op.
ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
if (failed(convert_embedded_lookup.VerifySignature())) {
return signalPassFailure();
}
convert_embedded_lookup.RewriteFunc();
} else if (attr.getValue() == mlir::TFL::kKerasLstm) {
func.eraseBody();
func.addEntryBlock();
OpBuilder builder(func.getBody());
if (failed(ConvertKerasLSTMLayer(func, &builder))) {
return signalPassFailure();
}
} else if (.....) /* Other fusions can plug in here */
}
Este é um snippet de código que mostra o mapeamento dessa operação composta para um objeto operação no LiteRT, aproveitando a função como uma interface de conversão.
void RewriteFunc() {
Value lookup = func_.getArgument(1);
Value value = func_.getArgument(0);
auto output_type = func_.getType().getResult(0);
OpBuilder builder(func_.getBody());
auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
func_.getLoc(), output_type, lookup, value);
builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
}