Использование компиляции AOT

Оптимизируйте свои подборки Сохраняйте и классифицируйте контент в соответствии со своими настройками.

Что такое tfcompile?

tfcompile представляет собой автономный инструмент , который дальновидность от времени (АОТ) компилирует TensorFlow графики в исполняемом коде. Это может уменьшить общий размер двоичного файла, а также избежать некоторых накладных расходов во время выполнения. Типичное использование-случай tfcompile должны составить график логического вывода в исполняемый код для мобильных устройств.

График TensorFlow обычно выполняется средой выполнения TensorFlow. Это влечет за собой некоторые накладные расходы времени выполнения для выполнения каждого узла в графе. Это также приводит к увеличению общего двоичного размера, поскольку в дополнение к самому графику должен быть доступен код для среды выполнения TensorFlow. Исполняемый код , созданный с помощью tfcompile не использует TensorFlow выполнения, и только имеет зависимость от ядер , которые фактически используются в вычислениях.

Компилятор построен на платформе XLA. Код моста TensorFlow к XLA каркасных пребывает под tensorflow / компилятор .

Что делает tfcompile?

tfcompile принимает подграф, обозначенное TensorFlow понятий кормов и выборок, и генерирует функцию , которая реализует , что подграф. Эти feeds являются входными аргументами для функции, и fetches являются выходными аргументами для функции. Все входные данные должны быть полностью указаны в фидах; Результирующий сокращенный подграф не может содержать узлы-заполнители или переменные. Обычно в качестве каналов указываются все заполнители и переменные, что гарантирует отсутствие этих узлов в результирующем подграфе. Генерируется функция упакованы как cc_library , с файлом заголовка экспортера функции, подпись и объектный файл , содержащий реализацию. Пользователь пишет код для вызова сгенерированной функции при необходимости.

Использование tfcompile

В этом разделе подробно шаги высокого уровня для генерации исполняемого двоичного с tfcompile из TensorFlow подграфа. Шаги следующие:

  • Шаг 1. Настройте подграф для компиляции
  • Шаг 2: Используйте tf_library сборки макроса для компиляции подграфа
  • Шаг 3. Напишите код для вызова подграфа
  • Шаг 4: Создайте окончательный двоичный файл

Шаг 1. Настройте подграф для компиляции

Определите каналы и выборки, которые соответствуют аргументам ввода и вывода для сгенерированной функции. Затем настройте feeds и fetches в 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" }
}

Шаг 2. Используйте макрос сборки tf_library для компиляции подграфа

Этот шаг преобразует график в cc_library с использованием tf_library сборки макрокоманды. cc_library состоит из объектного файла , содержащего код , сгенерированный из графика, вместе с файлом заголовок , который дает доступ к сгенерированному коду. tf_library использует tfcompile для компиляции TensorFlow графа в исполняемый код.

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

Для создания GraphDef прото (test_graph_tfmatmul.pb) для этого примера, выполнение make_test_graphs.py и указать расположение вывода с --out_dir флагом.

Типичные графики содержат Variables , представляющие весовые коэффициенты, которые выучили через тренировку, но tfcompile не может скомпилировать подграф , которые содержат Variables . Freeze_graph.py инструмент преобразует переменные в константы, используя значения , сохраненные в файле контрольных точек. Как удобства tf_library макрос поддерживает freeze_checkpoint аргумент, который запускает инструмент. Дополнительные примеры см tensorflow / компилятор / АОТ / тесты / BUILD .

Константы, которые отображаются в скомпилированном подграфе, компилируются непосредственно в сгенерированный код. Чтобы передать константы в сгенерированную функцию, а не компилировать их, просто передайте их как каналы.

Для получения дополнительной информации о tf_library сборки макроса см tfcompile.bzl .

Для получения дополнительной информации о подстилающей tfcompile инструмента, см tfcompile_main.cc .

Шаг 3. Напишите код для вызова подграфа

Этот шаг использует файл заголовка ( test_graph_tfmatmul.h ) , порожденную tf_library построения макроса в предыдущем шаге , чтобы ссылаться на сгенерированный код. Заголовочный файл расположен в bazel-bin каталог , соответствующий пакет сборки, и назван на основе набора атрибутов имя на tf_library построения макроса. Например, заголовок генерируется для test_graph_tfmatmul будет test_graph_tfmatmul.h . Ниже приводится сокращенная версия того, что сгенерировано. Созданный файл, в bazel-bin , содержит дополнительные полезные комментарии.

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

Сгенерированный класс C ++ называется MatMulComp в foo::bar имен, потому что это был cpp_class указан в tf_library макроса. Все сгенерированные классы имеют аналогичный API, с той лишь разницей, что методы обработки буферов аргументов и результатов. Эти методы различаются в зависимости от количества и типов буферов, которые были указаны в feed и fetch аргументы tf_library макрокоманды.

Существуют три типа буферов , управляемых в сгенерированном классе: args , представляющий входные данные, results , представляющие выходы, и temps , представляющий временные буфера , используемые внутри для выполнения вычислений. По умолчанию каждый экземпляр сгенерированного класса выделяет все эти буферы и управляет ими за вас. AllocMode конструктор аргумент может быть использован , чтобы изменить это поведение. Все буферы выровнены по границам 64 байта.

Сгенерированный класс C ++ - это просто оболочка низкоуровневого кода, сгенерированного XLA.

Пример вызова функции сгенерированной на основе 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;
}

Шаг 4: Создайте окончательный двоичный файл

Этот шаг сочетает в себе библиотеку , порожденную tf_library на шаге 2 , и код , написанный на шаге 3 , чтобы создать окончательный двоичный файл. Ниже приведен пример 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",
    ]
)