Cos'è tfcompile?
tfcompile è uno strumento autonomo che ahead-of-time (AOT) compilazioni tensorflow grafici in codice eseguibile. Può ridurre la dimensione binaria totale e anche evitare alcuni overhead di runtime. Un tipico caso d'uso di tfcompile è compilare un grafico inferenza in codice eseguibile per i dispositivi mobili.
Il grafico TensorFlow viene normalmente eseguito dal runtime di TensorFlow. Ciò comporta un sovraccarico di runtime per l'esecuzione di ciascun nodo nel grafico. Ciò porta anche a una dimensione binaria totale maggiore, poiché il codice per il runtime di TensorFlow deve essere disponibile, oltre al grafico stesso. Il codice eseguibile prodotto da tfcompile non utilizza il runtime tensorflow, e ha solo dipendenze kernel che vengono effettivamente utilizzati nel calcolo.
Il compilatore è costruito sulla base del framework XLA. Il codice colmare tensorflow alle risiede quadro XLA sotto tensorflow / compilatore .
Cosa fa tfcompile?
tfcompile prende un sottografo, identificata dai concetti tensorflow di feed e recuperi, e genera una funzione che implementa che sottografo. Le feeds sono gli argomenti di input per la funzione, ei fetches sono gli argomenti di uscita per la funzione. Tutti gli input devono essere completamente specificati dai feed; il sottografo tagliato risultante non può contenere nodi Segnaposto o Variabile. È comune specificare tutti i segnaposto e le variabili come feed, il che garantisce che il sottografo risultante non contenga più questi nodi. La funzione generata è confezionato come cc_library , con un file di intestazione esportazione della firma funzione, e un file oggetto contenente l'attuazione. L'utente scrive il codice per richiamare la funzione generata in modo appropriato.
Usando tfcompile
Questa sezione descrive passaggi di alto livello per la generazione di un binario eseguibile con tfcompile da un sottografo tensorflow. I passaggi sono:
- Passaggio 1: configurare il sottografo da compilare
- Fase 2: Utilizzare la
tf_librarymacro build per compilare il sottografo - Passaggio 3: scrivi il codice per richiamare il sottografo
- Passaggio 4: creare il file binario finale
Passaggio 1: configurare il sottografo da compilare
Identificare i feed e i recuperi che corrispondono agli argomenti di input e output per la funzione generata. Poi configurare le feeds e fetches in 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" }
}
Passaggio 2: utilizzare la macro build tf_library per compilare il sottografo
Questo passaggio converte il grafico in un cc_library usando l' tf_library macro accumulo. La cc_library costituito da un file oggetto contenente il codice generato dal grafico, insieme a un file di intestazione che dà accesso al codice generato. tf_library UTILIZZA tfcompile per compilare il grafico tensorflow in codice eseguibile.
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",
)
Per generare il proto GraphDef (test_graph_tfmatmul.pb) per questo esempio, eseguire make_test_graphs.py e specificare il percorso di uscita con la bandiera --out_dir.
Grafici tipici contengono Variables che rappresentano i pesi che vengono apprese attraverso la formazione, ma tfcompile non possono compilare un sottografo che contengono Variables . Il freeze_graph.py strumento converte variabili in costanti, utilizzando i valori memorizzati in un file di checkpoint. Per comodità, il tf_library macro sostiene il freeze_checkpoint argomento, che esegue lo strumento. Per ulteriori esempi vedere tensorflow / compilatore / AOT / test / build .
Le costanti che vengono visualizzate nel sottografo compilato vengono compilate direttamente nel codice generato. Per passare le costanti alla funzione generata, invece di averle compilate, passale semplicemente come feed.
Per i dettagli sulla tf_library macro di compilazione, vedere tfcompile.bzl .
Per dettagli sul sottostante tfcompile strumento, vedere tfcompile_main.cc .
Passaggio 3: scrivi il codice per richiamare il sottografo
Questo passaggio utilizza il file di intestazione ( test_graph_tfmatmul.h ) generato dal tf_library macro accumulo nel passaggio precedente per richiamare il codice generato. Il file di intestazione si trova nella bazel-bin directory corrispondente al pacchetto di build, e prende il nome in base al set di attributi nome sulla tf_library macro build. Ad esempio, l'intestazione generato per test_graph_tfmatmul sarebbe test_graph_tfmatmul.h . Di seguito è riportata una versione abbreviata di ciò che viene generato. Il file generato, in bazel-bin , contiene ulteriori commenti utili.
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 classe generata C ++ è chiamato MatMulComp nella foo::bar spazio dei nomi, perché quello era il cpp_class specificato nel tf_library macro. Tutte le classi generate hanno un'API simile, con l'unica differenza che sono i metodi per gestire arg e buffer dei risultati. Tali metodi diversi a seconda del numero e tipo di buffer, che sono stati specificati dal feed e fetch argomenti al tf_library macro.
Ci sono tre tipi di buffer gestite all'interno della classe generata: args rappresentano gli ingressi, results rappresentano le uscite e le temps rappresentano buffer temporanei utilizzati internamente per eseguire il calcolo. Per impostazione predefinita, ogni istanza della classe generata alloca e gestisce tutti questi buffer per te. AllocMode argomento del costruttore può essere usata per modificare questo comportamento. Tutti i buffer sono allineati ai limiti di 64 byte.
La classe C++ generata è solo un wrapper per il codice di basso livello generato da XLA.
Esempio di invocare la funzione generata sulla base 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;
}
Passaggio 4: creare il file binario finale
Questa fase combina la libreria generata da tf_library nel passaggio 2 e il codice scritto nel passaggio 3 per creare un binario finale. Di seguito è riportato un esempio bazel file di costruire.
# 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",
]
)