Qu'est-ce que tfcompile ?
tfcompile est un outil autonome qui compile en avance les graphiques TensorFlow en code exécutable. Cela peut réduire la taille binaire totale et éviter certains coûts d'exécution. Un cas d'utilisation typique de tfcompile consiste à compiler un graphique d'inférence en code exécutable pour les appareils mobiles.
Le graphe TensorFlow est normalement exécuté par l'environnement d'exécution TensorFlow. Cela entraîne une surcharge d'exécution liée à l'exécution de chaque nœud du graphe. Cela se traduit également par une taille binaire totale plus importante, car le code de l'environnement d'exécution TensorFlow doit être disponible, en plus du graphe lui-même. Le code exécutable généré par tfcompile n'utilise pas l'environnement d'exécution TensorFlow et ne possède des dépendances qu'avec des noyaux qui sont réellement utilisés dans le calcul.
Le compilateur repose sur le framework XLA. Le code reliant TensorFlow au framework XLA se trouve sous tensorflow/compiler.
À quoi sert tfcompile ?
tfcompile prend un sous-graphe identifié par les concepts de flux et d'extraction de TensorFlow, et génère une fonction qui l'implémente.
feeds sont les arguments d'entrée de la fonction et fetches sont les arguments de sortie de la fonction. Toutes les entrées doivent être entièrement spécifiées par les flux. Le sous-graphique élancé qui en résulte ne peut pas contenir de nœuds d'espace réservé ou de variable. Il est courant de spécifier tous les espaces réservés et toutes les variables en tant que flux, ce qui garantit que le sous-graphique résultant ne contient plus ces nœuds. La fonction générée est empaquetée sous la forme d'un fichier cc_library, avec un fichier d'en-tête exportant la signature de la fonction et un fichier d'objet contenant l'implémentation. L'utilisateur écrit du code pour appeler la fonction générée, le cas échéant.
Utiliser tfcompile
Cette section détaille les grandes étapes à suivre pour générer un binaire exécutable avec tfcompile à partir d'un sous-graphe TensorFlow. Voici la procédure à suivre :
- Étape 1: Configurer le sous-graphique à compiler
- Étape 2: Utiliser la macro de compilation
tf_librarypour compiler le sous-graphique - Étape 3: Écrire le code pour appeler le sous-graphique
- Étape 4: Créez le binaire final
Étape 1: Configurer le sous-graphique à compiler
Identifiez les flux et les extractions correspondant aux arguments d'entrée et de sortie de la fonction générée. Configurez ensuite feeds et fetches dans un fichier 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" }
}
Étape 2: Utiliser la macro de compilation tf_library pour compiler le sous-graphique
Cette étape convertit le graphique en cc_library à l'aide de la macro de compilation tf_library. cc_library se compose d'un fichier d'objet contenant le code généré à partir du graphique, ainsi que d'un fichier d'en-tête donnant accès au code généré. tf_library utilise tfcompile pour compiler le graphe TensorFlow en code exécutable.
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",
)
Pour générer le proto GraphDef (test_graph_tfmatmul.pb) pour cet exemple, exécutez make_test_graphs.py et spécifiez l'emplacement de sortie avec l'indicateur --out_dir.
Les graphiques classiques contiennent des Variables qui représentent les pondérations apprises via l'entraînement, mais tfcompile ne peut pas compiler un sous-graphe contenant Variables. L'outil freeze_graph.py convertit les variables en constantes à l'aide des valeurs stockées dans un fichier de point de contrôle. Pour plus de commodité, la macro tf_library prend en charge l'argument freeze_checkpoint, qui exécute l'outil. Pour plus d'exemples, consultez tensorflow/compiler/aot/tests/BUILD.
Les constantes qui s'affichent dans le sous-graphe compilé sont compilées directement dans le code généré. Pour transmettre les constantes à la fonction générée au lieu de les avoir compilées, il vous suffit de les transmettre en tant que flux.
Pour en savoir plus sur la macro de compilation tf_library, consultez tfcompile.bzl.
Pour en savoir plus sur l'outil tfcompile sous-jacent, consultez tfcompile_main.cc.
Étape 3: Écrire le code pour appeler le sous-graphique
Cette étape utilise le fichier d'en-tête (test_graph_tfmatmul.h) généré par la macro de compilation tf_library à l'étape précédente pour appeler le code généré. Le fichier d'en-tête se trouve dans le répertoire bazel-bin correspondant au package de compilation. Il est nommé en fonction de l'attribut de nom défini dans la macro de compilation tf_library. Par exemple, l'en-tête généré pour test_graph_tfmatmul serait test_graph_tfmatmul.h. Vous trouverez ci-dessous une version abrégée des éléments générés. Le fichier généré, dans bazel-bin, contient des commentaires utiles supplémentaires.
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 C++ générée est appelée MatMulComp dans l'espace de noms foo::bar, car il s'agit du cpp_class spécifié dans la macro tf_library. Toutes les classes générées ont une API similaire, la seule différence réside dans les méthodes de gestion des tampons d'argument et de résultat. Ces méthodes diffèrent en fonction du nombre et des types de tampons spécifiés par les arguments feed et fetch de la macro tf_library.
Trois types de tampons sont gérés dans la classe générée: args représente les entrées, results représente les sorties et temps représente les tampons temporaires utilisés en interne pour effectuer le calcul. Par défaut, chaque instance de la classe générée alloue et gère tous ces tampons à votre place. L'argument de constructeur AllocMode peut être utilisé pour modifier ce comportement. Tous les tampons sont alignés sur les limites de 64 octets.
La classe C++ générée n'est qu'un wrapper autour du code de bas niveau généré par XLA.
Exemple d'appel de la fonction générée en fonction de 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;
}
Étape 4: Créez le binaire final
Cette étape combine la bibliothèque générée par tf_library à l'étape 2 et le code écrit à l'étape 3 pour créer un binaire final. Vous trouverez ci-dessous un exemple de fichier BUILD bazel.
# 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",
]
)