Visão geral
Nesta página, descrevemos o design e as etapas necessárias para converter operações compostas no TensorFlow em operações fundidas no TensorFlow Lite. Essa infraestrutura é de uso geral e é compatível com a conversão de qualquer operação composta no TensorFlow para uma operação fundida correspondente no TensorFlow Lite.
Um exemplo de uso dessa infraestrutura é a fusão de operações de RNN do TensorFlow com o TensorFlow Lite, conforme detalhado aqui.
O que são operações fundidas
As operações do TensorFlow podem ser operações primitivas, como tf.add, ou podem ser compostas de outras operações primitivas, como tf.einsum. Uma operação primitiva aparece como um único nó no gráfico do TensorFlow, enquanto uma operação composta é uma coleção de nós no gráfico do TensorFlow. A execução de uma operação composta é equivalente à execução de cada uma das operações primitivas que as compõem.
Uma operação fundida corresponde a uma única operação que inclui toda a computação realizada por cada operação primitiva na operação composta correspondente.
Benefícios das operações combinadas
Existem operações combinadas para maximizar o desempenho das implementações do kernel subjacente, otimizando a computação geral e reduzindo o consumo de memória. Isso é muito valioso, especialmente para cargas de trabalho de inferência de baixa latência e plataformas móveis com recursos limitados.
As operações fundidas também fornecem uma interface de nível superior para definir transformações complexas, como a quantização, o que seria inviável ou muito difícil de fazer em um nível mais granular.
O TensorFlow Lite tem muitas instâncias de operações fundidas pelos motivos articulados acima. Essas operações combinadas normalmente correspondem a operações compostas no programa de origem do TensorFlow. Os exemplos de operações compostas no TensorFlow que são implementadas como uma única operação fundida no TensorFlow Lite incluem várias operações de RNN, como sequência unidirecional e bidirecional LSTM, conv2d (conv2d, bias add, relu), totalmente conectadas (matmul, bias add, relu) e muito mais. No TensorFlow Lite, a quantização LSTM é atualmente implementada apenas nas operações LSTM combinadas.
Desafios com operações combinadas
Converter operações compostas do TensorFlow em operações fundidas no TensorFlow Lite é 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 desafiador identificar (por exemplo, por meio da correspondência de padrões) o subgráfico correspondente a uma operação composta.
É possível que haja mais de uma implementação do TensorFlow direcionada a uma operação combinada do TensorFlow Lite. Por exemplo, há muitas implementações de LSTM no TensorFlow (Keras, Babelfish/lingvo etc.), e cada uma delas é composta por diferentes operações primitivas. No entanto, todas elas ainda podem ser convertidas para a mesma operação de LSTM combinada no TensorFlow Lite.
Por isso, a conversão de operações combinadas tem se mostrado bastante desafiadora.
Como converter da op composta para uma operação personalizada do TFLite (recomendado)
Una a operação composta em um tf.function.
Em muitos casos, parte do modelo pode ser mapeada para uma única operação no
TFLite. Isso pode ajudar na performance ao escrever uma implementação otimizada
para operações específicas. Para criar uma operação combinada no TFLite,
identifique a parte do gráfico que representa essa operação e a una em
um tf.function com
o atributo "experimental_implements" para um tf.function, que tem o valor de
atributo tfl_fusable_op com o valor true. Se a operação personalizada receber
atributos, transmita-os como parte das mesmas "experimental_implements".
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, porque
o atributo tfl_fusable_op já indica isso.
Implementar operação personalizada e registrar com o TFLite Interpreter
Implemente a operação combinada como uma operação personalizada do TFLite. Consulte as instructions.
O nome de registro da operação precisa ser semelhante ao nome
especificado no atributo name na assinatura das implementações.
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 uma operação composta para fundida (avançado)
A arquitetura geral para converter operações compostas do TensorFlow em operações fundidas do TensorFlow Lite é mostrada abaixo:
Una a operação composta em um tf.function.
No código-fonte do modelo do TensorFlow, identifique e abstraia a operação composta em uma tf.function com a anotação de função experimental_implements. Veja um exemplo de pesquisa de incorporação. A função define a interface e os argumentos dela precisam ser usados para implementar a lógica de conversão.
Escrever código de conversão
O código de conversão é escrito pela interface da função com a anotação implements. Veja um exemplo de fusão para pesquisa
de incorporação. Conceitualmente, o código de conversão substitui a implementação composta
dessa interface pela fundida.
Na passagem prepare-composite-functions, use o plug-in no código de conversão.
Em usos mais avançados, é possível implementar transformações complexas dos operados da operação composta para derivar os operandos da operação combinada. Veja o código de conversão Keras LSTM como exemplo.
Converter para o TensorFlow Lite
Use a API TFLiteConverter.from_saved_model para converter para o TensorFlow Lite.
funcionamento interno
Agora, descrevemos detalhes de alto nível do projeto geral na conversão para operações fundidas no TensorFlow Lite.
Como criar operações no TensorFlow
O uso de tf.function com o atributo de função experimental_implements permite que os usuários componham explicitamente novas operações usando operações primitivas do TensorFlow e especifiquem a interface que a operação composta resultante implementa. Isso é muito útil porque oferece:
- Um limite bem definido para a operação composta no gráfico subjacente do TensorFlow.
- Especifique explicitamente a interface implementada por essa operação. Os argumentos da tf.function correspondem aos argumentos desta interface.
Por exemplo, vamos considerar uma operação composta definida para implementar a pesquisa de incorporação. Ele é mapeado para uma operação combinada no TensorFlow Lite.
@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
Ao fazer com que os modelos usem operações compostas via tf.function, como ilustrado acima, é possível criar uma infraestrutura geral para identificar e converter essas operações em operações do TensorFlow Lite.
Como ampliar o conversor do TensorFlow Lite
O conversor do TensorFlow Lite lançado no início deste ano só era compatível com a importação de modelos do TensorFlow como um gráfico com todas as variáveis substituídas pelos valores de constantes correspondentes. Isso não funciona na fusão de operações, já que esses gráficos têm todas as funções inline para que as variáveis possam ser transformadas em constantes.
Para aproveitar a função tf.function com o recurso experimental_implements durante o processo de conversão, as funções precisam ser preservadas até mais tarde no processo de conversão.
Dessa forma, implementamos um novo fluxo de trabalho de importação e conversão de modelos do TensorFlow 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 o MLIR
- operações compostas de fusão
- análise de mutabilidade de variáveis
- congelar todas as variáveis somente leitura
Isso nos permite realizar a fusão de operações usando as funções que representam as operações compostas antes da função em linha e congelamento variável.
Como implementar a fusão de operações
Vamos conhecer melhor a passagem de fusão de operações. Esse cartão faz o seguinte:
- Repita todas as funções no módulo MLIR.
- Se uma função tiver o atributo tf._implements, com base no valor do atributo, o utilitário de fusão de operações apropriado será chamado.
- O utilitário de fusão de operações opera nos operandos e atributos da função (que servem como a interface para a conversão) e substitui o corpo da função por um corpo de função equivalente contendo a operação fundida.
- Em muitos casos, o corpo substituído vai conter operações diferentes da operação combinada. Elas correspondem a algumas transformações estáticas nos operadores da função para receber os operandos da operação combinada. Como esses cálculos podem ser sempre dobrados, eles não estariam presentes no flatbuffer exportado, em que apenas a operação fundida existiria.
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 é o snippet de código que mostra o mapeamento dessa operação composta para uma operação musa no TensorFlow Lite, usando 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());
}