¿Qué es tfcompile?
tfcompile es una herramienta independiente que compila gráficos de TensorFlow por adelantado (AOT) en código ejecutable. Puede reducir el tamaño total del objeto binario y también evitar algunas sobrecargas del entorno de ejecución. Un caso de uso típico de tfcompile es compilar un gráfico de inferencia en código ejecutable para dispositivos móviles.
El entorno de ejecución de TensorFlow suele ejecutar el grafo de TensorFlow. Esto genera cierta sobrecarga del entorno de ejecución por la ejecución de cada nodo en el grafo. Esto también conduce a un tamaño binario total mayor, ya que el código para el entorno de ejecución de TensorFlow debe estar disponible, además del grafo en sí. El código ejecutable que produce tfcompile no usa el entorno de ejecución de TensorFlow y solo tiene dependencias en los kernels que se usan en el procesamiento.
El compilador se compila sobre el framework de XLA. El código que conecta TensorFlow con el framework de XLA reside en tensorflow/compiler.
¿Qué hace tfcompile?
tfcompile toma un subgrafo, identificado por los conceptos de TensorFlow sobre feeds y recuperaciones, y genera una función que implementa ese subgrafo.
feeds son los argumentos de entrada de la función y fetches son los argumentos de salida de la función. Los feeds deben especificar todas las entradas por completo; el subgrafo resultante reducido no puede contener nodos de marcadores de posición ni variables. Es común especificar todos los marcadores de posición y las variables como feeds, lo que garantiza que el subgrafo resultante ya no contenga estos nodos. La función generada se empaqueta como cc_library, con un archivo de encabezado que exporta la firma de la función y un archivo de objeto que contiene la implementación. El usuario escribe código para invocar la función generada según corresponda.
Usa tfcompile
En esta sección, se detallan los pasos de alto nivel para generar un objeto binario ejecutable con tfcompile desde un subgrafo de TensorFlow. Estos son los pasos:
- Paso 1: Configura el subgrafo para compilar
- Paso 2: Usa la macro de compilación
tf_librarypara compilar el subgrafo - Paso 3: Escribe el código para invocar el subgrafo
- Paso 4: Crea el objeto binario final
Paso 1: Configura el subgrafo para compilar
Identifica los feeds y las recuperaciones que corresponden a los argumentos de entrada y salida de la función generada. Luego, configura feeds y fetches en un archivo proto tensorflow.tf2xla.Config.
# 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: Usa la macro de compilación tf_library para compilar el subgrafo
En este paso, se convierte el gráfico en un cc_library con la macro de compilación tf_library. El cc_library consiste en un archivo de objeto que contiene el código generado a partir del gráfico, junto con un archivo de encabezado que brinda acceso al código generado. tf_library usa tfcompile para compilar el grafo de 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 de GraphDef (test_graph_tfmatmul.pb) para este ejemplo, ejecuta make_test_graphs.py y especifica la ubicación de salida con la marca --out_dir.
Los gráficos típicos contienen Variables, que representa las ponderaciones que se aprenden a través del entrenamiento, pero tfcompile no puede compilar un subgrafo que contenga Variables. La herramienta freeze_graph.py convierte variables en constantes mediante valores almacenados en un archivo de punto de control. Para tu conveniencia, la macro tf_library admite el argumento freeze_checkpoint, que ejecuta la herramienta. Para ver más ejemplos, consulta tensorflow/compiler/aot/tests/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 compilarlas, simplemente pásalas como feeds.
Para obtener detalles sobre la macro de compilación tf_library, consulta tfcompile.bzl.
Para obtener detalles sobre la herramienta tfcompile subyacente, consulta tfcompile_main.cc.
Paso 3: Escribe el código para invocar el subgrafo
En este paso, se usa el archivo de encabezado (test_graph_tfmatmul.h) que generó la macro de compilación tf_library en el paso anterior para invocar el código generado. El archivo de encabezado se encuentra en el directorio bazel-bin correspondiente al paquete de compilación y se llama según el atributo de nombre configurado en la macro de compilación tf_library. Por ejemplo, el encabezado generado 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 útiles adicionales.
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 de C++ generada se llama MatMulComp en el espacio de nombres foo::bar, ya que ese era el cpp_class especificado en la macro tf_library. Todas las clases generadas tienen una API similar, con la única diferencia de los métodos para controlar los búferes de argumento y resultados. Esos métodos difieren según la cantidad y los tipos de búferes, que se especificaron mediante los argumentos feed y fetch de la macro tf_library.
Existen tres tipos de búferes administrados dentro de la clase generada: args, que representa las entradas, results, que representa las salidas, y temps, que representa los búferes temporales que se usan internamente para realizar el cálculo. De forma predeterminada, cada instancia de la clase generada asigna y administra todos estos búferes por ti. Se puede usar el argumento del constructor AllocMode para cambiar este comportamiento. Todos los búferes se alinean con límites de 64 bytes.
La clase de C++ generada es solo un wrapper del código de bajo nivel que genera XLA.
Ejemplo de invocación de la función generada basada 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 "third_party/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 objeto binario final
En este paso, se combina la biblioteca que genera tf_library en el paso 2 con el código escrito en el paso 3 para crear un objeto binario final. A continuación, se muestra un ejemplo de archivo bazel BUILD.
# 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",
]
)