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 LiteRT. Essa infraestrutura é de uso geral e oferece suporte à conversão de qualquer operação combinada no TensorFlow para uma operação de fusão correspondente no LiteRT.
Um exemplo de uso dessa infraestrutura é a fusão de operações de RNN do TensorFlow com o LiteRT, conforme detalhado aqui.
O que são operações de fusão?
As operações do TensorFlow podem ser primitivas, por exemplo, tf.add, ou compostas de outras operações primitivas, por exemplo, 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. Executar uma operação composta é equivalente a executar cada uma das operações primitivas constituintes.
Uma operação fundida corresponde a uma única operação que engloba todo o cálculo realizado por cada operação primitiva na operação composta correspondente.
Benefícios das operações fundidas
As operações de fusão existem para maximizar o desempenho das implementações de kernel subjacentes, 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 oferecem uma interface de nível mais alto para definir transformações complexas, como quantização, que seriam inviáveis ou muito difíceis de fazer em um nível mais granular.
A LiteRT tem muitas instâncias de operações fundidas pelos motivos articulados acima. Essas operações combinadas normalmente correspondem a operações compostas no programa TensorFlow de origem. Exemplos de operações compostas no TensorFlow que são implementadas como uma única operação fundida no LiteRT incluem várias operações de RNN, como LSTM de sequência unidirecional e bidirecional, convolução (conv2d, adição de bias, relu), totalmente conectada (matmul, adição de bias, relu) e muito mais. No LiteRT, a quantização de LSTM só é implementada nas operações de LSTM fundidas.
Desafios com operações fundidas
Converter operações compostas do TensorFlow para operações fundidas no 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 identificar (por exemplo, por correspondência de padrões) o subgrafo correspondente a uma operação composta.
Pode haver mais de uma implementação do TensorFlow segmentando uma operação LiteRT fundida. Por exemplo, há muitas implementações de LSTM no TensorFlow (Keras, Babelfish/lingvo etc.), e cada uma delas é composta de diferentes operações primitivas, mas todas ainda podem ser convertidas na mesma operação de LSTM fundida no LiteRT.
Por isso, a conversão de operações fundidas tem se mostrado bastante desafiadora.
Como converter uma operação composta em uma operação personalizada do TFLite (recomendado)
Envolva 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 no desempenho ao escrever uma implementação otimizada
para operações específicas. Para criar uma operação fundida no TFLite, identifique a parte do gráfico que representa uma operação fundida e encapsule-a em uma tf.function com o atributo "experimental_implements" para um tf.function, que tem o valor do atributo tfl_fusable_op com o valor true. Se a operação personalizada usar atributos, transmita-os como parte do mesmo "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, já que o atributo
tfl_fusable_op já implica isso.
Implementar uma operação personalizada e registrar com o interpretador do TFLite
Implemente sua operação combinada como uma operação personalizada do TFLite. Consulte as instruções.
O nome para registrar a 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 uma operação composta para uma fundida (avançado)
A arquitetura geral para converter operações compostas do TensorFlow em operações fundidas do LiteRT está abaixo:
Envolva 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. Confira um exemplo de pesquisa de incorporação. A função define a interface, e os argumentos dela 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 a anotação implements. Confira 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 combinada.
Na transmissão prepare-composite-functions, adicione seu código de conversão.
Em usos mais avançados, é possível implementar transformações complexas dos operandos da operação composta para derivar os operandos da operação fundida. Consulte o código de conversão do Keras LSTM (em inglês) como exemplo.
Converter para LiteRT
Use a API TFLiteConverter.from_saved_model para converter em LiteRT.
Configurações avançadas
Agora vamos descrever detalhes de alto nível do design geral na conversão para operações combinadas no LiteRT.
Como compor operações no TensorFlow
O uso de tf.function com o atributo de função experimental_implements permite que os usuários criem 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 do TensorFlow subjacente.
- Especifique explicitamente a interface que esta operação implementa. Os argumentos de tf.function correspondem aos argumentos dessa interface.
Por exemplo, considere uma operação composta definida para implementar a pesquisa de incorporação. Isso é mapeado para uma operação combinada 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
Ao fazer com que os modelos usem operações compostas via tf.function, conforme ilustrado acima, é possível criar uma infraestrutura geral para identificar e converter essas operações em operações LiteRT combinadas.
Como estender o conversor LiteRT
O conversor LiteRT lançado no início deste ano só permitia importar modelos do TensorFlow como um gráfico com todas as variáveis substituídas pelos valores constantes correspondentes. Isso não funciona para a 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 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.
Por isso, 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:
- Importar modelos salvos do TensorFlow para o MLIR
- fuse composite operations
- análise de mutabilidade de variáveis
- congelar todas as variáveis somente leitura
Isso permite que façamos a fusão de operações usando as funções que representam as operações compostas antes da incorporação de funções e do congelamento de variáveis.
Implementar a fusão de operações
Vamos analisar a transmissão de fusão de operações com mais detalhes. Essa transmissão faz o seguinte:
- Faça um loop em todas as funções do módulo MLIR.
- Se uma função tiver o atributo tf._implements, com base no valor do atributo, ela vai chamar o utilitário de fusão de operações apropriado.
- O utilitário de fusão de operações opera nos operandos e atributos da função (que servem como interface para a conversão) e substitui o corpo da função por um corpo de função equivalente que contém a operação fundida.
- Em muitos casos, o corpo substituído vai conter operações diferentes da operação de fusão. Elas correspondem a algumas transformações estáticas nos operandos da função para obter os operandos da operação combinada. Como todos esses cálculos podem ser constantes, eles não estariam presentes no flatbuffer exportado, em que apenas a operação combinada existiria.
Confira um snippet de código da transmissã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 */
}
Confira um snippet de código que mostra o mapeamento dessa operação composta para uma operação combinada no LiteRT 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());
}