¡Google I/O es una envoltura! Póngase al día con las sesiones de TensorFlow Ver sesiones

Usando la compilación AOT

¿Qué es tfcompile?

tfcompile es una herramienta independiente que antes-de-tiempo (AOT) compila TensorFlow gráficos en código ejecutable. Puede reducir el tamaño binario total y también evitar algunos gastos generales en tiempo de ejecución. Un caso de uso típico de tfcompile es compilar un gráfico de la inferencia en código ejecutable para dispositivos móviles.

El gráfico de TensorFlow normalmente lo ejecuta el tiempo de ejecución de TensorFlow. Esto genera una sobrecarga de tiempo de ejecución para la ejecución de cada nodo en el gráfico. Esto también conduce a un tamaño binario total más grande, ya que el código para el tiempo de ejecución de TensorFlow debe estar disponible, además del gráfico en sí. El código ejecutable producido por tfcompile no utilizar el tiempo de ejecución TensorFlow, y sólo tiene dependencias en los núcleos que realmente se utilizan en el cálculo.

El compilador está construido sobre el marco XLA. El código de puente a TensorFlow reside el marco XLA bajo tensorflow / compilador .

¿Qué hace tfcompile?

tfcompile toma un subgrafo, identificado por los conceptos TensorFlow de piensos y recuperaciones, y genera una función que implementos que subgrafo. Las feeds son los argumentos de entrada para la función, y las fetches son los argumentos de salida para la función. Todas las entradas deben estar completamente especificadas por las fuentes; el subgrafo podado resultante no puede contener nodos de marcador de posición o variable. Es común especificar todos los marcadores de posición y variables como fuentes, lo que garantiza que el subgráfico resultante ya no contenga estos nodos. La función generada se empaqueta como un cc_library , con un archivo de cabecera exportar la firma de la función, y un fichero objeto que contiene la aplicación. El usuario escribe código para invocar la función generada según corresponda.

Usando tfcompile

Esta sección detalla los pasos de alto nivel para la generación de un binario ejecutable con tfcompile de un subgrafo TensorFlow. Los pasos son:

  • Paso 1: configurar el subgrafo para compilar
  • Paso 2: Utilice el tf_library acumulación macro para compilar el subgrafo
  • Paso 3: escriba el código para invocar el subgrafo
  • Paso 4: crea el binario final

Paso 1: configurar el subgrafo para compilar

Identifique los feeds y recuperaciones que corresponden a los argumentos de entrada y salida de la función generada. A continuación, configure los feeds y fetches en un tensorflow.tf2xla.Config proto.

# Each feed is a positional input argument for the generated function.  The order
# of each entry matches the order of each input argument.  Here “x_hold” and “y_hold”
# refer to the names of placeholder nodes defined in the graph.
feed {
  id { node_name: "x_hold" }
  shape {
    dim { size: 2 }
    dim { size: 3 }
  }
}
feed {
  id { node_name: "y_hold" }
  shape {
    dim { size: 3 }
    dim { size: 2 }
  }
}

# Each fetch is a positional output argument for the generated function.  The order
# of each entry matches the order of each output argument.  Here “x_y_prod”
# refers to the name of a matmul node defined in the graph.
fetch {
  id { node_name: "x_y_prod" }
}

Paso 2: use la macro de compilación tf_library para compilar el subgrafo

Esta etapa convierte la gráfica en un cc_library usando el tf_library macro de construcción. El cc_library consta de un fichero objeto que contiene el código generado a partir del gráfico, junto con un archivo de cabecera que da acceso al código generado. tf_library utiliza tfcompile para compilar el gráfico TensorFlow en código ejecutable.

load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# Use the tf_library macro to compile your graph into executable code.
tf_library(
    # name is used to generate the following underlying build rules:
    # <name>           : cc_library packaging the generated header and object files
    # <name>_test      : cc_test containing a simple test and benchmark
    # <name>_benchmark : cc_binary containing a stand-alone benchmark with minimal deps;
    #                    can be run on a mobile device
    name = "test_graph_tfmatmul",
    # cpp_class specifies the name of the generated C++ class, with namespaces allowed.
    # The class will be generated in the given namespace(s), or if no namespaces are
    # given, within the global namespace.
    cpp_class = "foo::bar::MatMulComp",
    # graph is the input GraphDef proto, by default expected in binary format.  To
    # use the text format instead, just use the ‘.pbtxt’ suffix.  A subgraph will be
    # created from this input graph, with feeds as inputs and fetches as outputs.
    # No Placeholder or Variable ops may exist in this subgraph.
    graph = "test_graph_tfmatmul.pb",
    # config is the input Config proto, by default expected in binary format.  To
    # use the text format instead, use the ‘.pbtxt’ suffix.  This is where the
    # feeds and fetches were specified above, in the previous step.
    config = "test_graph_tfmatmul.config.pbtxt",
)

Para generar el proto GraphDef (test_graph_tfmatmul.pb) para este ejemplo, ejecute make_test_graphs.py y especifique la ubicación de salida con la bandera --out_dir.

Gráficos típicos contienen Variables que representan los pesos que se aprenden a través de la formación, pero tfcompile no puede compilar un subgrafo que contienen Variables . El freeze_graph.py herramienta convierte variables en constantes, utilizando los valores almacenados en un archivo de controles. Para su conveniencia, la tf_library macro apoya la freeze_checkpoint argumento, que se ejecuta la herramienta. Para ver más ejemplos tensorflow / compilador / AOT / pruebas / BUILD .

Las constantes que aparecen en el subgrafo compilado se compilan directamente en el código generado. Para pasar las constantes a la función generada, en lugar de tenerlas compiladas, simplemente pásalas como fuentes.

Para más detalles sobre la tf_library macro de construcción, ver tfcompile.bzl .

Para más detalles sobre el subyacente tfcompile herramienta, consulte tfcompile_main.cc .

Paso 3: escriba el código para invocar el subgrafo

Este paso se utiliza el archivo de cabecera ( test_graph_tfmatmul.h ) generada por el tf_library macro acumulación en el paso anterior para invocar el código generado. El archivo de cabecera se encuentra en el bazel-bin directorio correspondiente a la construcción del paquete, y lleva el nombre basado en el conjunto nombre de atributo en el tf_library macro de construcción. Por ejemplo, la cabecera generada para test_graph_tfmatmul sería test_graph_tfmatmul.h . A continuación se muestra una versión abreviada de lo que se genera. El archivo generado, en bazel-bin , contiene comentarios adicionales útiles.

namespace foo {
namespace bar {

// MatMulComp represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code.
class MatMulComp {
 public:
  // AllocMode controls the buffer allocation mode.
  enum class AllocMode {
    ARGS_RESULTS_AND_TEMPS,  // Allocate arg, result and temp buffers
    RESULTS_AND_TEMPS_ONLY,  // Only allocate result and temp buffers
  };

  MatMulComp(AllocMode mode = AllocMode::ARGS_RESULTS_AND_TEMPS);
  ~MatMulComp();

  // Runs the computation, with inputs read from arg buffers, and outputs
  // written to result buffers. Returns true on success and false on failure.
  bool Run();

  // Arg methods for managing input buffers. Buffers are in row-major order.
  // There is a set of methods for each positional argument.
  void** args();

  void set_arg0_data(float* data);
  float* arg0_data();
  float& arg0(size_t dim0, size_t dim1);

  void set_arg1_data(float* data);
  float* arg1_data();
  float& arg1(size_t dim0, size_t dim1);

  // Result methods for managing output buffers. Buffers are in row-major order.
  // Must only be called after a successful Run call. There is a set of methods
  // for each positional result.
  void** results();


  float* result0_data();
  float& result0(size_t dim0, size_t dim1);
};

}  // end namespace bar
}  // end namespace foo

La clase generada C ++ se llama MatMulComp en el foo::bar espacio de nombres, porque esa era la cpp_class especificado en el tf_library macro. Todas las clases generadas tienen una API similar, con la única diferencia de los métodos para manejar arg y búferes de resultados. Estos métodos difieren según el número y tipos de los tampones, que fueron especificados por la feed y fetch argumentos a la tf_library macro.

Hay tres tipos de tampones administrados dentro de la clase generada: args que representan los insumos, results representan las salidas, y temps que representan buffers temporales de uso interno para realizar el cálculo. De forma predeterminada, cada instancia de la clase generada asigna y administra todos estos búferes por usted. El AllocMode argumento del constructor puede ser utilizado para cambiar este comportamiento. Todos los búferes están alineados con límites de 64 bytes.

La clase C ++ generada es solo una envoltura del código de bajo nivel generado por XLA.

Ejemplo de invocar la función generada basado en tfcompile_test.cc :

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include <iostream>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated

int main(int argc, char** argv) {
  Eigen::ThreadPool tp(2);  // Size the thread pool as appropriate.
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());


  foo::bar::MatMulComp matmul;
  matmul.set_thread_pool(&device);

  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, matmul.arg0_data());
  std::copy(args + 6, args + 12, matmul.arg1_data());
  matmul.Run();

  // Check result
  if (matmul.result0(0, 0) == 58) {
    std::cout << "Success" << std::endl;
  } else {
    std::cout << "Failed. Expected value 58 at 0,0. Got:"
              << matmul.result0(0, 0) << std::endl;
  }

  return 0;
}

Paso 4: crea el binario final

Este paso combina la biblioteca generada por tf_library en el paso 2 y el código escrito en el paso 3 para crear un binario final. A continuación se muestra un ejemplo bazel fichero de construcción.

# Example of linking your binary
# Also see //tensorflow/compiler/aot/tests/BUILD
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# The same tf_library call from step 2 above.
tf_library(
    name = "test_graph_tfmatmul",
    ...
)

# The executable code generated by tf_library can then be linked into your code.
cc_binary(
    name = "my_binary",
    srcs = [
        "my_code.cc",  # include test_graph_tfmatmul.h to access the generated header
    ],
    deps = [
        ":test_graph_tfmatmul",  # link in the generated object file
        "//third_party/eigen3",
    ],
    linkopts = [
          "-lpthread",
    ]
)