Ayuda a proteger la Gran Barrera de Coral con TensorFlow en Kaggle Únete Challenge

Fusión de operaciones de TensorFlow

Descripción general

En esta página, se describen el diseño y los pasos necesarios para convertir operaciones compuestas en TensorFlow en operaciones fusionadas en TensorFlow Lite. Esta infraestructura es de uso 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 operaciones de 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 a partir de 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. La ejecución de una operación compuesta equivale a la ejecución de cada una de sus operaciones primitivas constituyentes.

Una operación fusionada corresponde a una sola operación que incluye todo el cálculo realizado por cada operación primitiva dentro de la operación compuesta correspondiente.

Beneficios de las operaciones fusionadas

Las operaciones fusionadas existen para maximizar el rendimiento de sus implementaciones de kernel subyacentes, optimizando el cálculo general y reduciendo la huella 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 a un nivel más granular.

TensorFlow Lite tiene muchas instancias de operaciones fusionadas por las razones expuestas anteriormente. Estas operaciones fusionadas suelen corresponder a operaciones compuestas en el programa TensorFlow de origen. Los ejemplos de operaciones compuestas en TensorFlow que se implementan como una sola operación fusionada en TensorFlow Lite incluyen varias operaciones RNN como LSTM de secuencia unidireccional y bidireccional, convolución (conv2d, bias add, relu), completamente conectado (matmul, bias add, relu) y más . En TensorFlow Lite, la cuantificación de LSTM actualmente solo se implementa en las operaciones de LSTM fusionadas.

Desafíos con operaciones fusionadas

Convertir operaciones compuestas de TensorFlow en 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 difícil identificar (por ejemplo, mediante la 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 todas podrían convertirse a la misma operación LSTM fusionada en TensorFlow Lite.

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

Conversión de funcionamiento compuesto a fusionado

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

dibujo

Envuelva la operación compuesta en una función tf.function

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

Escribe el 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 de implements . Vea un ejemplo de fusión para incrustar búsqueda . Conceptualmente, el código de conversión reemplaza la implementación compuesta de esta interfaz con la fusionada.

En el paso prepare-composite-functions, 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.

Convierta a TensorFlow Lite

Usa 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 mediante 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 de TensorFlow subyacente.
  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 de incrustaciones. 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 tales operaciones en operaciones fusionadas de TensorFlow Lite.

Ampliación del convertidor de TensorFlow Lite

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

Para aprovechar la función 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. Específicamente, las nuevas características agregadas son:

  1. Importar modelos guardados de TensorFlow en MLIR
  2. operaciones compuestas de fusibles
  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 del alineamiento de funciones y la congelación de variables.

Implementación de fusión de operaciones

Veamos la operación de paso 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, basado en el valor del atributo, llama a la utilidad de fusión de operación adecuada.
  3. La utilidad de fusión de operaciones opera en 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 de 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 se pueden plegar constantemente, 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 */
}

A continuación, se muestra un fragmento de código que muestra la asignación de esta operación compuesta a una operación fusionada en TensorFlow Lite que aprovecha 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());
  }