Diese Seite wurde von der Cloud Translation API übersetzt.
Switch to English

Verwenden der AOT-Kompilierung

Was ist tfcompile?

tfcompile ist ein eigenständiges Tool, mit dem AOT (TensorFlow-Diagramme im Voraus) in ausführbaren Code kompilieren. Dies kann die Gesamtgröße der Binärdateien reduzieren und auch einige Laufzeit-Overheads vermeiden. Ein typischer Anwendungsfall von tfcompile ist das Kompilieren eines Inferenzgraphen in ausführbaren Code für mobile Geräte.

Das TensorFlow-Diagramm wird normalerweise von der TensorFlow-Laufzeit ausgeführt. Dies verursacht einen gewissen Laufzeitaufwand für die Ausführung jedes Knotens im Diagramm. Dies führt auch zu einer größeren Gesamtbinärgröße, da der Code für die TensorFlow-Laufzeit zusätzlich zum Diagramm selbst verfügbar sein muss. Der von tfcompile erzeugte ausführbare Code verwendet nicht die TensorFlow-Laufzeit und ist nur von Kerneln abhängig, die tatsächlich für die Berechnung verwendet werden.

Der Compiler basiert auf dem XLA-Framework. Der Code, der TensorFlow mit dem XLA-Framework verbindet, befindet sich unter Tensorflow / Compiler .

Was macht tfcompile?

tfcompile einen Untergraphen, der durch die TensorFlow-Konzepte für Feeds und Abrufe identifiziert wird, und generiert eine Funktion, die diesen Untergraphen implementiert. Die feeds sind die Eingabeargumente für die Funktion, und die fetches sind die Ausgabeargumente für die Funktion. Alle Eingänge müssen vollständig von den Feeds spezifiziert werden. Der resultierende beschnittene Untergraph darf keine Platzhalter- oder Variablenknoten enthalten. Es ist üblich, alle Platzhalter und Variablen als Feeds anzugeben, um sicherzustellen, dass der resultierende Untergraph diese Knoten nicht mehr enthält. Die generierte Funktion wird als cc_library , wobei eine Header-Datei die Funktionssignatur exportiert und eine Objektdatei die Implementierung enthält. Der Benutzer schreibt Code, um die generierte Funktion entsprechend aufzurufen.

Verwenden von tfcompile

In diesem Abschnitt werden tfcompile Schritte zum Generieren einer ausführbaren Binärdatei mit tfcompile aus einem TensorFlow-Untergraphen beschrieben. Die Schritte sind:

  • Schritt 1: Konfigurieren Sie den zu kompilierenden Untergraphen
  • Schritt 2: Verwenden Sie das Build-Makro tf_library , um den Untergraphen zu kompilieren
  • Schritt 3: Schreiben Sie Code, um den Untergraphen aufzurufen
  • Schritt 4: Erstellen Sie die endgültige Binärdatei

Schritt 1: Konfigurieren Sie den zu kompilierenden Untergraphen

Identifizieren Sie die Feeds und Abrufe, die den Eingabe- und Ausgabeargumenten für die generierte Funktion entsprechen. Konfigurieren Sie dann die feeds und fetches in einem 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" }
}
 

Schritt 2: Verwenden Sie das Build-Makro tf_library, um den Untergraphen zu kompilieren

Dieser Schritt konvertiert das Diagramm mithilfe des Build-Makros tf_library in eine cc_library . Die cc_library besteht aus einer Objektdatei, die den aus dem Diagramm generierten Code enthält, sowie einer Header-Datei, die den Zugriff auf den generierten Code ermöglicht. tf_library verwendet tfcompile , um das TensorFlow-Diagramm in ausführbaren Code zu kompilieren.

 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",
)
 

Um das GraphDef-Proto (test_graph_tfmatmul.pb) für dieses Beispiel zu generieren, führen Sie make_test_graphs.py aus und geben Sie den Ausgabeort mit dem Flag --out_dir an.

Typische Diagramme enthalten Variables die die Gewichte darstellen, die durch Training gelernt wurden. tfcompile kann jedoch keinen Untergraphen kompilieren, der Variables enthält. Das Tool freeze_graph.py konvertiert Variablen in Konstanten unter Verwendung von Werten, die in einer Prüfpunktdatei gespeichert sind. Der tf_library unterstützt das Makro freeze_checkpoint Argument freeze_checkpoint , mit dem das Tool ausgeführt wird. Weitere Beispiele finden Sie unter tensorflow / compiler / aot / tests / BUILD .

Konstanten, die im kompilierten Untergraphen angezeigt werden, werden direkt in den generierten Code kompiliert. Um die Konstanten an die generierte Funktion zu übergeben, anstatt sie kompilieren zu lassen, übergeben Sie sie einfach als Feeds.

Einzelheiten zur tf_library Build Makro finden tfcompile.bzl .

Einzelheiten zum zugrunde liegenden tfcompile Tool finden Sie unter tfcompile_main.cc .

Schritt 3: Schreiben Sie Code, um den Untergraphen aufzurufen

In diesem Schritt wird die vom Build-Makro tf_library im vorherigen Schritt generierte Header-Datei ( test_graph_tfmatmul.h ) verwendet, um den generierten Code aufzurufen. Die Header-Datei befindet sich im bazel-bin , das dem Build-Paket entspricht, und wird basierend auf dem im Build-Makro tf_library festgelegten Attribut tf_library . Der für test_graph_tfmatmul generierte test_graph_tfmatmul wäre beispielsweise test_graph_tfmatmul.h . Unten finden Sie eine Kurzfassung dessen, was generiert wird. Die generierte Datei in bazel-bin enthält zusätzliche nützliche Kommentare.

 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
 

Die generierte C ++ Klasse heißt MatMulComp in den foo::bar Namespace, denn das ist die war cpp_class im angegebenen tf_library Makro. Alle generierten Klassen haben eine ähnliche API. Der einzige Unterschied besteht in den Methoden zum Umgang mit Arg- und Ergebnispuffern. Diese Methoden unterscheiden sich je nach Anzahl und Typ der Puffer, die durch die Argumente feed und fetch für das Makro tf_library .

Innerhalb der generierten Klasse werden drei Arten von Puffern verwaltet: args für die Eingaben, results für die Ausgaben und temps für temporäre Puffer, die intern zur Durchführung der Berechnung verwendet werden. Standardmäßig weist jede Instanz der generierten Klasse alle diese Puffer für Sie zu und verwaltet sie. Das AllocMode Konstruktorargument kann verwendet werden, um dieses Verhalten zu ändern. Alle Puffer sind an 64-Byte-Grenzen ausgerichtet.

Die generierte C ++ - Klasse ist nur ein Wrapper um den von XLA generierten Low-Level-Code.

Beispiel für den Aufruf der generierten Funktion basierend auf 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;
}
 

Schritt 4: Erstellen Sie die endgültige Binärdatei

Dieser Schritt kombiniert die von tf_library in Schritt 2 generierte Bibliothek und den in Schritt 3 geschriebenen Code, um eine endgültige Binärdatei zu erstellen. Unten finden Sie eine Beispieldatei für 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",
    ]
)