Fusão de operação do TensorFlow

Visão geral

Esta página descreve 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 oferece suporte à conversão de qualquer operação composta no TensorFlow para uma operação combinada correspondente no TensorFlow Lite.

Um exemplo de uso dessa infraestrutura é a fusão da operação do TensorFlow RNN com o TensorFlow Lite, conforme detalhado aqui .

O que são operações fundidas

desenho

As operações do TensorFlow podem ser operações primitivas, por exemplo, tf.add , ou podem ser 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 de suas operações primitivas constituintes.

Uma operação fundida corresponde a uma única operação que inclui toda a computação realizada por cada operação primitiva dentro da operação composta correspondente.

Benefícios das operações fundidas

As operações fundidas existem para maximizar o desempenho de suas 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 restrição de recursos.

As operações fundidas também fornecem uma interface de nível superior para definir transformações complexas como quantização, que de outra forma seriam inviáveis ​​ou muito difíceis 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 fundidas geralmente 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 TensorFlow Lite incluem várias operações RNN, como sequência unidirecional e bidirecional LSTM, convolução (conv2d, bias add, relu), totalmente conectada (matmul, bias add, relu) e muito mais . No TensorFlow Lite, a quantização LSTM atualmente é implementada apenas nas operações LSTM fundidas.

Desafios com operações fundidas

Converter operações compostas do TensorFlow em operações fundidas no TensorFlow Lite é um problema difícil. Isto é porque:

  1. 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 correspondência de padrões) o subgráfico correspondente a tal operação composta.

  2. Pode haver mais de uma implementação do TensorFlow visando uma operação fundida do TensorFlow Lite. Por exemplo, existem 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 para a mesma operação LSTM fundida no TensorFlow Lite.

Como tal, a conversão de operações fundidas provou ser bastante desafiadora.

Envolva a operação composta em um tf.function

Em muitos casos, alguma 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 poder criar uma operação fundida em TFLite, identifique a parte do gráfico que representa uma operação fundida e envolva-a 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, passe-os como parte dos mesmos "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))

Observe que você não precisa definir allow_custom_ops no conversor, pois o atributo tfl_fusable_op já implica isso.

Implemente a operação personalizada e registre-se no TFLite Interpreter

Implemente sua operação fundida como uma operação TFLite Custom - veja as instruções .

Observe que o nome para registrar o op deve ser semelhante ao nome especificado no atributo name na assinatura do implemento.

Um exemplo para o op 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, &reg);

Conversão de operação composta para operação fundida (Avançado)

A arquitetura geral para converter operações compostas do TensorFlow em operações fundidas do TensorFlow Lite está abaixo:

desenho

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 um tf.function com a anotação da função experimental_implements . Veja um exemplo de pesquisa de incorporação . A função 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 pela interface da função com a anotação implements . Veja um exemplo de fusão para incorporar pesquisa . Conceitualmente, o código de conversão substitui a implementação composta dessa interface pela integrada.

No passo prepare-composite-functions, insira 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 Keras LSTM . código de conversão como um exemplo.

Converter para TensorFlow Lite

Use a API TFLiteConverter.from_saved_model para converter para o TensorFlow Lite.

Sob o capô

Agora descrevemos detalhes de alto nível do design geral na conversão para operações fundidas no TensorFlow Lite.

Compondo 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, pois fornece:

  1. Um limite bem definido para a operação composta no gráfico subjacente do TensorFlow.
  2. Especifique explicitamente a interface que esta operação implementa. Os argumentos da função tf.function correspondem aos argumentos desta interface.

Como exemplo, vamos considerar uma operação composta definida para implementar a pesquisa de incorporação. Isso é mapeado para uma operação fundida 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 por meio de tf.function conforme ilustrado acima, torna-se possível construir uma infraestrutura geral para identificar e converter essas operações em operações fundidas do TensorFlow Lite.

Estendendo o conversor TensorFlow Lite

O conversor TensorFlow Lite que foi lançado no início deste ano suportava apenas a importação de modelos TensorFlow como um gráfico com todas as variáveis ​​substituídas por seus valores constantes correspondentes. Isso não funciona para a operação de fusão, pois esses gráficos têm todas as funções embutidas para que as variáveis ​​possam ser transformadas em constantes.

Para alavancar 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.

Como tal, implementamos um novo fluxo de trabalho de importação e conversão de modelos do TensorFlow no conversor para dar suporte ao caso de uso de fusão de operação composta. Especificamente, os novos recursos adicionados são:

  1. Importando modelos salvos do TensorFlow para MLIR
  2. operações compostas de fusíveis
  3. análise de mutabilidade variável
  4. 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 inlining e do congelamento de variáveis.

Implementando a fusão de operações

Vejamos o passo de fusão da operação com mais detalhes. Este passe faz o seguinte:

  1. Faça um loop por todas as funções no módulo MLIR.
  2. Se uma função tiver o atributo tf._implements, com base no valor do atributo, chama o utilitário de fusão de operação apropriado.
  3. O utilitário de fusão de operação 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 contendo a operação de fusão.
  4. Em muitos casos, o corpo substituído conterá outras operações além da operação com fusível. Estes correspondem a algumas transformações estáticas nos operandos da função para obter os operandos da operação fundida. Como esses cálculos podem ser dobrados constantemente, eles não estariam presentes no flatbuffer exportado onde existiria apenas a operação fundida.

Aqui está o trecho de código da passagem mostrando 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 */
}

Aqui está o snippet de código que mostra o mapeamento dessa operação composta para uma operação fundida no TensorFlow Lite, 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());
  }