Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

Korzystanie z kompilacji AOT

Co to jest tfcompile?

tfcompile to samodzielne narzędzie, które z wyprzedzeniem (AOT) kompiluje wykresy TensorFlow do postaci kodu wykonywalnego. Może zmniejszyć całkowity rozmiar plików binarnych, a także uniknąć niektórych kosztów ogólnych w czasie wykonywania. Typowym przypadkiem tfcompile jest skompilowanie wykresu wnioskowania do kodu wykonywalnego dla urządzeń mobilnych.

Wykres TensorFlow jest zwykle wykonywany przez środowisko wykonawcze TensorFlow. Wiąże się to z pewnym narzutem czasu wykonywania na wykonanie każdego węzła na wykresie. Prowadzi to również do większego całkowitego rozmiaru binarnego, ponieważ oprócz samego wykresu musi być dostępny kod środowiska wykonawczego TensorFlow. Kod wykonywalny utworzony przez tfcompile nie korzysta ze środowiska wykonawczego TensorFlow i ma tylko zależności od jądra, które są faktycznie używane w obliczeniach.

Kompilator jest zbudowany na frameworku XLA. Kod łączący TensorFlow ze strukturą XLA znajduje się w tensorflow / compiler .

Co robi tfcompile?

tfcompile pobiera podgraf, identyfikowany przez koncepcje kanałów i pobierania TensorFlow, i generuje funkcję, która implementuje ten podgraf. feeds są argumentami wejściowymi funkcji, a fetches są argumentami wyjściowymi funkcji. Wszystkie dane wejściowe muszą być w pełni określone przez źródła danych; wynikowy przycięty podgraf nie może zawierać węzłów zastępczych ani zmiennych. Powszechnie określa się wszystkie symbole zastępcze i zmienne jako źródła danych, co zapewnia, że ​​wynikowy podgraf nie zawiera już tych węzłów. Wygenerowana funkcja jest spakowana jako cc_library , z plikiem nagłówkowym eksportującym sygnaturę funkcji i plikiem obiektu zawierającym implementację. Użytkownik pisze kod, aby w razie potrzeby wywołać wygenerowaną funkcję.

Korzystanie z tfcompile

Ta sekcja szczegółowo opisuje kroki wysokiego poziomu generowania wykonywalnego pliku binarnego za pomocą tfcompile z podgrafu TensorFlow. Kroki są następujące:

  • Krok 1: Skonfiguruj podgraf do kompilacji
  • Krok 2: Użyj makra kompilacji tf_library aby skompilować podgraf
  • Krok 3: Napisz kod wywołujący podgraf
  • Krok 4: Utwórz ostateczny plik binarny

Krok 1: Skonfiguruj podgraf do kompilacji

Zidentyfikuj źródła i pobrania, które odpowiadają argumentom wejściowym i wyjściowym dla wygenerowanej funkcji. Następnie skonfiguruj feeds i fetches w 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" }
}

Krok 2: Użyj makra kompilacji tf_library, aby skompilować podgraf

Ten krok konwertuje wykres do cc_library przy użyciu tf_library kompilacji tf_library . Biblioteka cc_library składa się z pliku obiektowego zawierającego kod wygenerowany na podstawie wykresu, wraz z plikiem nagłówkowym, który daje dostęp do wygenerowanego kodu. tf_library wykorzystuje tfcompile do kompilowania wykresu TensorFlow do kodu wykonywalnego.

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

Aby wygenerować protokół GraphDef (test_graph_tfmatmul.pb) dla tego przykładu, uruchom make_test_graphs.py i określ lokalizację wyjściową za pomocą flagi --out_dir.

Typowe wykresy zawierają Variables reprezentujące wagi, których uczy się podczas uczenia, ale tfcompile nie może skompilować podgrafu zawierającego Variables . Narzędzie freeze_graph.py konwertuje zmienne na stałe, używając wartości przechowywanych w pliku punktów kontrolnych. Dla wygody makro tf_library obsługuje argument freeze_checkpoint , który uruchamia narzędzie. Więcej przykładów znajdziesz w tensorflow / compiler / aot / tests / BUILD .

Stałe, które pojawiają się w skompilowanym podgrafie, są kompilowane bezpośrednio do wygenerowanego kodu. Aby przekazać stałe do wygenerowanej funkcji, zamiast ich wkompilowywania, po prostu przekaż je jako źródła.

Aby uzyskać szczegółowe informacje na temat tf_library kompilacji tf_library , zobacz tfcompile.bzl .

Aby uzyskać szczegółowe informacje na temat bazowego narzędzia tfcompile , zobacz tfcompile_main.cc .

Krok 3: Napisz kod wywołujący podgraf

W tym kroku do wywołania wygenerowanego kodu używany jest plik nagłówkowy ( test_graph_tfmatmul.h ) wygenerowany przez makro kompilacji tf_library w poprzednim kroku. Plik nagłówkowy znajduje się w katalogu bazel-bin odpowiadającym pakietowi kompilacji i jest nazywany na podstawie atrybutu name ustawionego w tf_library kompilacji tf_library . Na przykład nagłówek wygenerowany dla test_graph_tfmatmul miałby test_graph_tfmatmul.h . Poniżej znajduje się skrócona wersja tego, co jest generowane. Wygenerowany plik w bazel-bin zawiera dodatkowe przydatne komentarze.

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

Wygenerowana klasa C ++ nosi nazwę MatMulComp w przestrzeni nazw foo::bar , ponieważ była to cpp_class określona w makrze tf_library . Wszystkie wygenerowane klasy mają podobny interfejs API, a jedyną różnicą są metody do obsługi arg i buforów wyników. Sposoby te różnią się w zależności od ilości i rodzaju buforów, które zostały określone przez feed i fetch argumentów do tf_library makro.

Istnieją trzy rodzaje buforów zarządzanych w ramach wygenerowanej klasy: args reprezentujące wejść results reprezentujące wyjść i temps reprezentujące tymczasowych buforów używane wewnętrznie do wykonywania obliczeń. Domyślnie każde wystąpienie wygenerowanej klasy przydziela wszystkie te bufory i zarządza nimi za Ciebie. Do zmiany tego zachowania można użyć argumentu konstruktora AllocMode . Wszystkie bufory są wyrównane do granic 64-bajtowych.

Wygenerowana klasa C ++ jest tylko otoką wokół kodu niskiego poziomu generowanego przez XLA.

Przykład wywołania wygenerowanej funkcji na podstawie 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;
}

Krok 4: Utwórz ostateczny plik binarny

Ten krok łączy bibliotekę wygenerowaną przez tf_library w kroku 2 i kod napisany w kroku 3 w celu utworzenia ostatecznego pliku binarnego. Poniżej znajduje się przykładowy plik 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",
    ]
)