Diese Seite wurde von der Cloud Translation API übersetzt.

AOT-Kompilierung verwenden

Was ist tfcompile?

tfcompile ist ein eigenständiges Tool, das TensorFlow-Grafiken im Voraus (AOT) in ausführbaren Code kompiliert. Es kann die Gesamtgröße der Binärprogramme reduzieren und einige Laufzeit-Overheads vermeiden. Ein typischer Anwendungsfall von tfcompile ist die Kompilierung eines Inferenzdiagramms in ausführbarem Code für Mobilgeräte.

Die TensorFlow-Grafik 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 neben der Grafik selbst der Code für die TensorFlow-Laufzeit verfügbar sein muss. Der von tfcompile erstellte ausführbare Code verwendet nicht die TensorFlow-Laufzeit und weist nur Abhängigkeiten von Kerneln auf, 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 verwendet eine Teilgrafik, die durch die TensorFlow-Konzepte Feeds und Abrufe identifiziert wird, und generiert eine Funktion zur Implementierung dieser Teilgrafik. feeds sind die Eingabeargumente für die Funktion und fetches die Ausgabeargumente für die Funktion. Alle Eingaben müssen vollständig durch die Feeds angegeben werden. Das resultierende, bereinigte Teildiagramm darf keine Platzhalter- oder Variablenknoten enthalten. Es ist üblich, alle Platzhalter und Variablen als Feeds anzugeben. Dadurch wird sichergestellt, dass diese Knoten in der resultierenden Teilgrafik nicht mehr enthalten sind. Die generierte Funktion wird als cc_library gepackt, mit einer Headerdatei, die die Funktionssignatur exportiert, und einer Objektdatei, die die Implementierung enthält. Der Nutzer schreibt Code, um die generierte Funktion gegebenenfalls aufzurufen.

tfcompile verwenden

In diesem Abschnitt werden allgemeine Schritte zum Generieren einer ausführbaren Binärdatei mit tfcompile aus einer TensorFlow-Teilgrafik beschrieben. Folgende Schritte sind auszuführen:

Schritt 1: Zu kompilierende Teilgrafik konfigurieren
Schritt 2: Mit dem Build-Makro tf_library die Teilgrafik kompilieren
Schritt 3: Code zum Aufrufen der Teilgrafik schreiben
Schritt 4: Endgültige Binärdatei erstellen

Schritt 1: Zu kompilierende Teilgrafik konfigurieren

Identifizieren Sie die Feeds und Abrufe, die den Eingabe- und Ausgabeargumenten für die generierte Funktion entsprechen. Konfigurieren Sie dann 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: Mit dem Build-Makro „tf_library“ die Teilgrafik kompilieren

Dieser Schritt wandelt das Diagramm mithilfe des Build-Makros tf_library in ein cc_library um. cc_library besteht aus einer Objektdatei mit dem aus der Grafik generierten Code und einer Headerdatei, die Zugriff auf den generierten Code gewährt. tf_library verwendet tfcompile, um die TensorFlow-Grafik 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",
)

Führen Sie make_test_graphs.py aus und geben Sie den Ausgabespeicherort mit dem Flag "--out_dir" an, um das GraphDef-Protokoll (test_graph_tfmatmul.pb) für dieses Beispiel zu generieren.

Typische Grafiken enthalten Variables, das die Gewichtungen darstellt, die durch Training erlernt werden. tfcompile kann jedoch keine Teilgrafik mit Variables kompilieren. Das Tool freeze_graph.py wandelt Variablen in Konstanten um. Dazu werden Werte verwendet, die in einer Prüfpunktdatei gespeichert sind. Der Einfachheit halber unterstützt das Makro tf_library das Argument freeze_checkpoint, mit dem das Tool ausgeführt wird. Weitere Beispiele finden Sie unter tensorflow/compiler/aot/tests/BUILD.

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

Weitere Informationen zum Build-Makro tf_library finden Sie unter tfcompile.bzl.

Weitere Informationen zum zugrunde liegenden tfcompile-Tool finden Sie unter tfcompile_main.cc.

Schritt 3: Code zum Aufrufen der Teilgrafik schreiben

In diesem Schritt wird die Headerdatei (test_graph_tfmatmul.h) verwendet, die im vorherigen Schritt vom Build-Makro tf_library generiert wurde, um den generierten Code aufzurufen. Die Headerdatei befindet sich im Verzeichnis bazel-bin, das dem Build-Paket entspricht, und wird nach dem Attribut „name“ benannt, das im Build-Makro tf_library festgelegt ist. Der für test_graph_tfmatmul generierte Header wäre beispielsweise test_graph_tfmatmul.h. Unten sehen Sie eine gekürzte Version dessen, was generiert wird. Die generierte Datei in bazel-bin enthält weitere 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 im Namespace foo::bar, da dies der im Makro tf_library angegebene cpp_class war. Alle generierten Klassen haben eine ähnliche API, mit dem einzigen Unterschied, dass sie die Methoden zur Verarbeitung von Arg- und Ergebniszwischenspeichern haben. Diese Methoden unterscheiden sich je nach Anzahl und Typ der Zwischenspeicher, die mit den Argumenten feed und fetch im Makro tf_library angegeben wurden.

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 ordnet und verwaltet jede Instanz der generierten Klasse alle diese Puffer für Sie. Sie können dieses Verhalten mit dem Konstruktorargument AllocMode ä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 das Aufrufen der generierten Funktion anhand von 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;
}

Schritt 4: Endgültige Binärdatei erstellen

In diesem Schritt wird die in Schritt 2 von tf_library generierte Bibliothek mit dem in Schritt 3 geschriebenen Code kombiniert, um eine endgültige Binärdatei zu erstellen. Im Folgenden finden Sie ein Beispiel für eine bazel-Build-Datei.

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