Fusión de operación de TensorFlow

Descripción general

Esta página describe el diseño y los pasos necesarios para convertir operaciones compuestas en TensorFlow en operaciones fusionadas en TensorFlow Lite. Esta infraestructura es de propósito general y admite la conversión de cualquier operación compuesta en TensorFlow a una operación fusionada correspondiente en TensorFlow Lite.

Un ejemplo de uso de esta infraestructura es la fusión de la operación TensorFlow RNN con TensorFlow Lite, como se detalla aquí .

¿Qué son las operaciones fusionadas?

dibujo

Las operaciones de TensorFlow pueden ser operaciones primitivas, por ejemplo, tf.add , o pueden estar compuestas por otras operaciones primitivas, por ejemplo, tf.einsum . Una operación primitiva se muestra como un solo nodo en el gráfico de TensorFlow, mientras que una operación compuesta es una colección de nodos en el gráfico de TensorFlow. Ejecutar una operación compuesta equivale a ejecutar cada una de sus operaciones primitivas constituyentes.

Una operación fusionada corresponde a una única operación que incluye todos los cálculos realizados por cada operación primitiva dentro de la operación compuesta correspondiente.

Beneficios de las operaciones fusionadas

Existen operaciones fusionadas para maximizar el rendimiento de sus implementaciones de kernel subyacentes, optimizando el cálculo general y reduciendo el uso de memoria. Esto es muy valioso, especialmente para cargas de trabajo de inferencia de baja latencia y plataformas móviles con recursos limitados.

Las operaciones fusionadas también proporcionan una interfaz de nivel superior para definir transformaciones complejas como la cuantificación, que de otro modo serían inviables o muy difíciles de realizar en un nivel más granular.

TensorFlow Lite tiene muchos casos de operaciones fusionadas por las razones expuestas anteriormente. Estas operaciones fusionadas generalmente corresponden a operaciones compuestas en el programa fuente TensorFlow. Ejemplos de operaciones compuestas en TensorFlow que se implementan como una única operación fusionada en TensorFlow Lite incluyen varias operaciones RNN como secuencia unidireccional y bidireccional LSTM, convolución (conv2d, sesgo add, relu), completamente conectada (matmul, sesgo add, relu) y más. . En TensorFlow Lite, la cuantificación LSTM actualmente solo se implementa en las operaciones LSTM fusionadas.

Desafíos con operaciones fusionadas

Convertir operaciones compuestas de TensorFlow a operaciones fusionadas en TensorFlow Lite es un problema difícil. Esto es porque:

  1. Las operaciones compuestas se representan en el gráfico de TensorFlow como un conjunto de operaciones primitivas sin un límite bien definido. Puede resultar muy complicado identificar (por ejemplo, mediante coincidencia de patrones) el subgráfico correspondiente a dicha operación compuesta.

  2. Puede haber más de una implementación de TensorFlow dirigida a una operación fusionada de TensorFlow Lite. Por ejemplo, hay muchas implementaciones de LSTM en TensorFlow (Keras, Babelfish/lingvo, etc.) y cada una de ellas se compone de diferentes operaciones primitivas, pero aún así todas podrían convertirse a la misma operación LSTM fusionada en TensorFlow Lite.

Como tal, la conversión de operaciones fusionadas ha resultado todo un desafío.

Envuelva la operación compuesta en una tf.function

En muchos casos, alguna parte del modelo se puede asignar a una única operación en TFLite. Esto puede ayudar con el rendimiento al escribir una implementación optimizada para operaciones específicas. Para poder crear una operación fusionada en TFLite, identifique la parte del gráfico que representa una operación fusionada y envuélvala en un tf.function con el atributo "experimental_implements" en un tf.function , que tiene el valor del atributo tfl_fusable_op con valor true . Si la operación personalizada toma atributos, páselos como parte de los mismos "implementos_experimentales".

Ejemplo,

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))

Tenga en cuenta que no necesita configurar allow_custom_ops en el convertidor ya que el atributo tfl_fusable_op ya implica esto.

Implemente operaciones personalizadas y regístrese con TFLite Interpreter

Implemente su operación fusionada como una operación personalizada de TFLite; consulte las instrucciones .

Tenga en cuenta que el nombre con el que registrar la operación debe ser similar al nombre especificado en el atributo name en la firma del implemento.

Un ejemplo de la operación en el ejemplo es

  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);

Conversión de operación compuesta a fusionada (avanzada)

La arquitectura general para convertir operaciones compuestas de TensorFlow en operaciones fusionadas de TensorFlow Lite se muestra a continuación:

dibujo

Envuelva la operación compuesta en una tf.function

En el código fuente del modelo TensorFlow, identifique y resuma la operación compuesta en una tf.function con la anotación de función experimental_implements . Vea un ejemplo de búsqueda integrada . La función define la interfaz y sus argumentos deben usarse para implementar la lógica de conversión.

Escribir código de conversión

El código de conversión se escribe según la interfaz de la función con la anotación implements . Vea un ejemplo de fusión para incrustar búsquedas . Conceptualmente, el código de conversión reemplaza la implementación compuesta de esta interfaz por la fusionada.

En el paso de preparación de funciones compuestas, agregue su código de conversión .

En usos más avanzados, es posible implementar transformaciones complejas de los operandos de la operación compuesta para derivar los operandos de la operación fusionada. Consulte Keras LSTM . código de conversión como ejemplo.

Convertir a TensorFlow Lite

Utilice la API TFLiteConverter.from_saved_model para convertir a TensorFlow Lite.

Bajo el capó

Ahora describimos detalles de alto nivel del diseño general en la conversión a operaciones fusionadas en TensorFlow Lite.

Componer operaciones en TensorFlow

El uso de tf.function con el atributo de función experimental_implements permite a los usuarios componer explícitamente nuevas operaciones utilizando operaciones primitivas de TensorFlow y especificar la interfaz que implementa la operación compuesta resultante. Esto es muy útil ya que proporciona:

  1. Un límite bien definido para la operación compuesta en el gráfico subyacente de TensorFlow.
  2. Especifique explícitamente la interfaz que implementa esta operación. Los argumentos de tf.function corresponden a los argumentos de esta interfaz.

Como ejemplo, consideremos una operación compuesta definida para implementar la búsqueda incrustada. Esto se asigna a una operación fusionada en 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

Al hacer que los modelos utilicen operaciones compuestas a través de tf.function como se ilustra arriba, es posible construir una infraestructura general para identificar y convertir dichas operaciones en operaciones fusionadas de TensorFlow Lite.

Ampliando el convertidor TensorFlow Lite

El convertidor TensorFlow Lite que se lanzó a principios de este año solo admitía la importación de modelos TensorFlow como un gráfico con todas las variables reemplazadas por sus valores constantes correspondientes. Esto no funciona para la fusión de operaciones, ya que dichos gráficos tienen todas las funciones integradas para que las variables se puedan convertir en constantes.

Para aprovechar tf.function con la función experimental_implements durante el proceso de conversión, las funciones deben conservarse hasta más adelante en el proceso de conversión.

Como tal, implementamos un nuevo flujo de trabajo para importar y convertir modelos de TensorFlow en el convertidor para admitir el caso de uso de fusión de operaciones compuestas. En concreto, las nuevas características añadidas son:

  1. Importación de modelos guardados de TensorFlow a MLIR
  2. fusionar operaciones compuestas
  3. análisis de mutabilidad variable
  4. congelar todas las variables de solo lectura

Esto nos permite realizar la fusión de operaciones utilizando las funciones que representan las operaciones compuestas antes de la incorporación de funciones y la congelación de variables.

Implementación de la operación de fusión

Veamos la operación de pase de fusión con más detalle. Este pase hace lo siguiente:

  1. Recorra todas las funciones del módulo MLIR.
  2. Si una función tiene el atributo tf._implements, según el valor del atributo, llama a la utilidad de fusión de operaciones adecuada.
  3. La utilidad de fusión de operaciones opera sobre los operandos y atributos de la función (que sirven como interfaz para la conversión) y reemplaza el cuerpo de la función con un cuerpo de función equivalente que contiene la operación fusionada.
  4. En muchos casos, el cuerpo reemplazado contendrá operaciones distintas a la operación fusionada. Estos corresponden a algunas transformaciones estáticas en los operandos de la función para obtener los operandos de la operación fusionada. Dado que todos estos cálculos pueden plegarse de manera constante, no estarían presentes en el búfer plano exportado donde solo existiría la operación fusionada.

Aquí hay un fragmento de código del pase que muestra el flujo de trabajo 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 */
}

Aquí hay un fragmento de código que muestra la asignación de esta operación compuesta a una operación fusionada en TensorFlow Lite aprovechando la función como una interfaz de conversión.

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());
  }