Korzystanie z kompilacji AOT

Co to jest tf aktualnie?

tfcompile to samodzielne narzędzie, które z wyprzedzeniem kompiluje wykresy TensorFlow w kod wykonywalny. Może zmniejszyć całkowity rozmiar plików binarnych i uniknąć narzutu czasowego. Typowym przypadkiem użycia tfcompile jest kompilowanie grafu wnioskowania w kod wykonywalny na urządzenia mobilne.

Wykres TensorFlow jest zwykle wykonywany przez środowisko wykonawcze TensorFlow. Wiąże się to z nakładem pracy podczas wykonywania każdego węzła na wykresie. Prowadzi to również do większego całkowitego rozmiaru pliku 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ąder, które są rzeczywiście używane w obliczeniach.

Kompilator jest oparty na platformie XLA. Kod łączący TensorFlow z platformą XLA znajduje się w tensorflow/compiler.

Do czego służy tfbuild?

tfcompile pobiera podtekst określony przez koncepcje kanałów i pobrań w TensorFlow, a następnie generuje funkcję, która implementuje ten podgraf. feeds to argumenty wejściowe funkcji, a fetches to argumenty wyjściowe. Wszystkie dane wejściowe muszą być w pełni określone przez źródła. Wynikowy przycięty podgraf nie może zawierać węzłów zastępczych ani zmiennych. Zmienne te często określa się jako pliki danych, dzięki czemu wynikowy podpunkt nie będzie już zawierać tych węzłów. Wygenerowana funkcja ma postać pakietu cc_library – plik nagłówkowy eksportujący podpis funkcji oraz plik obiektu zawierający implementację. Użytkownik pisze kod, który wywołuje odpowiednią funkcję.

Korzystanie z tfbuild

W tej sekcji szczegółowo opisujemy czynności, jakie musisz wykonać, aby wygenerować wykonywalny plik binarny za pomocą funkcji tfcompile z podgrafu TensorFlow. Kroki:

  • Krok 1. Skonfiguruj podtytuł w celu skompilowania
  • Krok 2. Użyj makra kompilacji tf_library do skompilowania podgrafu
  • Krok 3. Napisz kod wywołujący podpunkt
  • Krok 4. Utwórz ostateczny plik binarny

Krok 1. Skonfiguruj podtytuł w celu skompilowania

Zidentyfikuj pliki danych i metody pobierania pasujące do argumentów wejściowych i wyjściowych dla wygenerowanej funkcji. Następnie skonfiguruj feeds i fetches w protokole 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 do skompilowania podgrafu.

Ten krok spowoduje przekształcenie wykresu w cc_library za pomocą makra kompilacji tf_library. cc_library składa się z pliku obiektu zawierającego kod wygenerowany z wykresu oraz pliku nagłówkowego, który daje dostęp do wygenerowanego kodu. tf_library wykorzystuje tfcompile do kompilowania grafu TensorFlow w kod wykonywalny.

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ć proto w GrafDef (test_graph_tfmatmul.pb) w tym przykładzie, uruchom polecenie make_test_graphs.py i określ lokalizację wyjściową za pomocą flagi --out_dir.

Typowe wykresy zawierają wartość Variables reprezentującą wagi nauczane podczas trenowania, ale parametr tfcompile nie może skompilować podtytułu zawierającego Variables. Narzędzie freeze_graph.py przekształca zmienne w stałe, używając wartości zapisanych 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 na stronie tensorflow/compiler/aot/tests/BUILD.

Stałe, które pojawiają się w skompilowanym podprogramie, są kompilowane bezpośrednio w wygenerowany kod. Aby przekazywać stałe do wygenerowanej funkcji, zamiast je skompilować, przekaż je jako kanały.

Szczegółowe informacje o makrze kompilacji tf_library znajdziesz w instrukcji tfcompile.bzl.

Szczegółowe informacje o podstawowym narzędziu tfcompile znajdziesz tutaj: tfcompile_main.cc.

Krok 3. Napisz kod wywołujący podpunkt

W tym kroku do wywołania wygenerowanego kodu używany jest plik nagłówka (test_graph_tfmatmul.h) wygenerowany przez makro kompilacji tf_library w poprzednim kroku. Plik nagłówka znajduje się w katalogu bazel-bin odpowiadającym pakietowi kompilacji i jego nazwa zależy od atrybutu nazwy ustawionego w makrze kompilacji tf_library. Na przykład nagłówek wygenerowany dla elementu test_graph_tfmatmul będzie wyglądał tak: test_graph_tfmatmul.h. Poniżej znajduje się skrócona wersja tego pliku. Wygenerowany plik w języku 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++ nazywa się MatMulComp w przestrzeni nazw foo::bar, ponieważ to jest klasa cpp_class określona w makrze tf_library. Wszystkie wygenerowane klasy mają podobny interfejs API. Jedyną różnicą są metody obsługi buforów argumentów i buforów wyników. Te metody różnią się w zależności od liczby i typów buforów, które zostały określone za pomocą argumentów feed i fetch makra tf_library.

W wygenerowanej klasie zarządzane są 3 typy buforów: args reprezentujące dane wejściowe, results reprezentujące dane wyjściowe oraz temps reprezentujące bufory tymczasowe używane wewnętrznie do wykonywania obliczeń. Domyślnie każda instancja wygenerowanej klasy przydziela wszystkie te bufory i zarządza nimi. Aby zmienić to działanie, można użyć argumentu konstruktora AllocMode. Wszystkie bufory są wyrównane względem 64-bajtowych obramowań.

Wygenerowana klasa C++ jest po prostu otoką kodu niskiego poziomu wygenerowanego 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 "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;
}

Krok 4. Utwórz ostateczny plik binarny

Ten krok łączy bibliotekę wygenerowaną przez tf_library w kroku 2 z kodem napisanym w kroku 3 w celu utworzenia ostatecznego pliku binarnego. Poniżej znajdziesz 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",
    ]
)