Cette page a été traduite par l'API Cloud Translation.

Utilisation de la compilation AOT

Qu'est-ce que tfcompile ?

tfcompile est un outil autonome qui en avance sur le temps (AOT) compile des graphiques tensorflow en code exécutable. Cela peut réduire la taille binaire totale et également éviter certains frais généraux d'exécution. Un cas d' utilisation typique de tfcompile consiste à compiler un graphe d'inférence en code exécutable pour les appareils mobiles.

Le graphique TensorFlow est normalement exécuté par l'environnement d'exécution TensorFlow. Cela entraîne une surcharge d'exécution pour l'exécution de chaque nœud du graphe. Cela conduit également à une taille binaire totale plus importante, car le code de l'environnement d'exécution TensorFlow doit être disponible, en plus du graphique lui-même. Le code exécutable produit par tfcompile ne pas utiliser le moteur d' exécution de tensorflow, et ne dispose que des dépendances sur les noyaux qui sont effectivement utilisés dans le calcul.

Le compilateur est construit sur le framework XLA. Le code de pontage à tensorflow réside le cadre dans le cadre XLA tensorflow / compilateur .

Que fait tfcompile ?

tfcompile prend un sous - graphe, identifié par les concepts de tensorflow des flux et récupérations, et génère une fonction qui implémente cette sous - graphe. Les feeds sont les arguments d'entrée pour la fonction et les 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-graphe élagué résultant 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 les variables en tant que flux, ce qui garantit que le sous-graphe résultant ne contient plus ces nœuds. La fonction générée est emballé sous forme de cc_library , avec un fichier d' en- tête d' exporter la signature de la fonction, et un fichier objet contenant l'application. L'utilisateur écrit du code pour appeler la fonction générée comme il convient.

Utiliser tfcompile

Cette section décrit les étapes de haut niveau pour générer un binaire exécutable avec tfcompile à partir d' un sous - graphe de tensorflow. Les étapes sont :

Étape 1 : Configurer le sous-graphe à compiler
Étape 2: Utilisez le tf_library macro build pour compiler le sous - graphe
Étape 3 : écrivez le code pour appeler le sous-graphe
Étape 4 : Créer le binaire final

Étape 1 : Configurer le sous-graphe à compiler

Identifiez les flux et les extractions qui correspondent aux arguments d'entrée et de sortie de la fonction générée. Configurez ensuite les feeds et les fetches dans 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" }
}

Étape 2 : utilisez la macro de génération tf_library pour compiler le sous-graphe

Cette étape convertit le graphique en cc_library utilisant la tf_library macro de construction. Le cc_library se compose d'un fichier objet contenant le code généré à partir du graphique, ainsi qu'un fichier d' en- tête qui donne accès au code généré. tf_library Utilise tfcompile pour compiler le graphique 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 (de test_graph_tfmatmul.pb) pour cet exemple, exécutez make_test_graphs.py et spécifiez l'emplacement de sortie avec le drapeau --out_dir.

Graphiques typiques contiennent des Variables représentant les poids qui sont apprises par la formation, mais tfcompile ne peut pas compiler un sous - graphe qui contiennent des Variables . Le freeze_graph.py outil convertit les variables en constantes, en utilisant les valeurs stockées dans un fichier de point de contrôle. Par commodité, la tf_library macro soutient le freeze_checkpoint argument, qui exécute l'outil. Pour plus d' exemples voir tensorflow / compilateur / AOT / tests / BUILD .

Les constantes qui apparaissent 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, plutôt que de les compiler, transmettez-les simplement en tant que flux.

Pour plus de détails sur la tf_library macro de construction, voir tfcompile.bzl .

Pour plus de détails sur le sous - jacent tfcompile outil, voir tfcompile_main.cc .

Étape 3 : écrivez le code pour appeler le sous-graphe

Cette étape utilise le fichier d' en- tête ( test_graph_tfmatmul.h ) généré par le tf_library macro de construction à l'étape précédente pour invoquer le code généré. Le fichier d' en- tête est situé dans le bazel-bin répertoire correspondant au package de construction, et est nommé en fonction de l'ensemble attribut name sur la tf_library macro de construction. 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 de ce qui est généré. Le fichier généré, dans bazel-bin , contient des commentaires supplémentaires utiles.

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é est appelé MatMulComp dans le foo::bar espace de noms, parce que ce fut le cpp_class spécifié dans la tf_library macro. Toutes les classes générées ont une API similaire, la seule différence étant les méthodes pour gérer les tampons d'argument et de résultat. Ces méthodes diffèrent en fonction du nombre et de types de tampons, qui ont été spécifiés par l' feed et fetch des arguments à la tf_library macro.

Il existe trois types de tampons gérés dans la classe générée: args représentant les entrées, les results représentant les sorties et les temps représentant des 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 pour vous. Le AllocMode argument du constructeur peut être utilisé pour modifier ce comportement. Tous les tampons sont alignés sur des 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'invoquer la fonction générée sur la base 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 "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éer 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 fichier binaire final. Voici un exemple bazel fichier 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",
    ]
)