Korzystanie z kompilacji AOT

Co to jest tfcompile?

tfcompile jest samodzielnym narzędziem, które z wyprzedzeniem-of-time (AOT) kompiluje TensorFlow wykresy do kodu wykonywalnego. Może zmniejszyć całkowity rozmiar pliku binarnego, a także uniknąć niektórych ogólnych kosztów działania. Typowym zastosowaniem-przypadek tfcompile jest opracowanie wykresu wnioskowania na kod wykonywalny dla urządzeń mobilnych.

Wykres TensorFlow jest zwykle wykonywany przez środowisko wykonawcze TensorFlow. Pociąga to za sobą pewne narzuty 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 produkowane przez tfcompile nie wykorzystuje środowisko wykonawcze TensorFlow i ma tylko współzależności na jądrach, które są faktycznie stosowane w obliczeniach.

Kompilator jest zbudowany na platformie XLA. Kod pomostowych TensorFlow do rezyduje ramowych XLA pod tensorflow / kompilator .

Co robi tfcompile?

tfcompile wykonuje subgraph określona przez koncepcje TensorFlow kanałów i pobiera oraz generuje funkcję narzędzi że podgrafu. Te feeds są argumenty wejściowe dla funkcji, a fetches są argumenty wyjściowe dla tej funkcji. Wszystkie dane wejściowe muszą być w pełni określone w kanałach; wynikowy przycięty podgraf nie może zawierać węzłów zastępczych ani zmiennych. Powszechne jest określanie wszystkich symboli zastępczych i zmiennych jako kanałów informacyjnych, co zapewnia, że ​​wynikowy podgraf nie zawiera już tych węzłów. Wygenerowany funkcja jest pakowany jako cc_library , z plikiem nagłówkowym wywozu podpis funkcji, a pliku obiektu zawierającego realizację. Użytkownik pisze kod, aby odpowiednio wywołać wygenerowaną funkcję.

Korzystanie z tfcompile

Ta sekcja Szczegóły High kroki poziomu generowania wykonywalnego binarnych z tfcompile z podgrafu TensorFlow. Kroki to:

  • Krok 1: Skonfiguruj podwykres do skompilowania
  • Krok 2: Użyj tf_library build makro skompilować podgrafu
  • Krok 3: Napisz kod, aby wywołać podgraf
  • Krok 4: Utwórz ostateczny plik binarny

Krok 1: Skonfiguruj podwykres do skompilowania

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

Krok 2: Użyj makra tf_library build do skompilowania podwykresu

Etap ten przekształca wykresu na cc_library używając tf_library gromadzeniu makro. cc_library składa się z obiektu pliku zawierającego kodu wygenerowanego z wykresu, wraz z plikiem nagłówka, który zapewnia dostęp do wygenerowanego kodu. tf_library wykorzystuje tfcompile do sporządzenia 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ć proto GraphDef (test_graph_tfmatmul.pb) dla tego przykład uruchomić make_test_graphs.py i określić położenie wyjściowe z --out_dir flagi.

Typowe wykresy zawierają Variables reprezentujące ciężarów, które są wyuczone poprzez szkolenia, ale tfcompile nie można skompilować podgrafu które zawierają Variables . Freeze_graph.py narzędzie konwertuje zmienne do stałych, przy użyciu wartości przechowywanych w pliku kontrolnego. Jako wygody tf_library makro obsługuje freeze_checkpoint argument, który uruchamia narzędzie. Do innych przykładów zobaczyć tensorflow / kompilatora / wyposażeniu m.in / testy / 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 kompilować je, po prostu przekaż je jako kanały.

Szczegółowe informacje na temat tf_library kompilacji makro, patrz tfcompile.bzl .

Szczegółowe informacje na temat podstawowej tfcompile narzędzia, zobacz tfcompile_main.cc .

Krok 3: Napisz kod, aby wywołać podgraf

Ten krok wykorzystuje plik nagłówka ( test_graph_tfmatmul.h ) wygenerowany przez tf_library kompilacji makro w poprzednim kroku do powoływania wygenerowany kod. Plik nagłówka znajduje się w bazel-bin katalogu odpowiadającego budowaniu pakietu i jest nazwany na podstawie zestawu atrybutów nazwa na tf_library kompilacji makro. Na przykład, nagłówek generowany dla test_graph_tfmatmul byłoby test_graph_tfmatmul.h . Poniżej znajduje się skrócona wersja tego, co jest generowane. Wygenerowany plik w bazel-bin , zawiera dodatkowe użyteczne 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

Wygenerowany klasa C ++ nazywa MatMulComp w foo::bar nazw, dlatego że był cpp_class określone w tf_library makro. Wszystkie wygenerowane klasy mają podobne API, a jedyną różnicą są metody do obsługi buforów arg i 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żda instancja wygenerowanej klasy przydziela i zarządza wszystkimi tymi buforami za Ciebie. AllocMode konstruktor argumentem może być wykorzystana do zmiany tego zachowania. Wszystkie bufory są wyrównane do granic 64-bajtowych.

Wygenerowana klasa C++ jest tylko opakowaniem wokół niskopoziomowego kodu generowanego przez XLA.

Przykład wywołanie funkcji generowane 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 etap łączy bibliotekę generowane przez tf_library w kroku 2 i kodu napisanego w kroku 3, aby utworzyć ostateczną binarny. Poniżej znajduje się przykład bazel plik budować.

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