Что такое tfcompile?
tfcompile — это автономный инструмент, который заранее (AOT) компилирует графики TensorFlow в исполняемый код. Это может уменьшить общий размер двоичного файла, а также избежать некоторых накладных расходов во время выполнения. Типичным вариантом использования tfcompile является компиляция графа вывода в исполняемый код для мобильных устройств.
Граф TensorFlow обычно выполняется средой выполнения TensorFlow. Это влечет за собой некоторые накладные расходы во время выполнения каждого узла графа. Это также приводит к увеличению общего двоичного размера, поскольку помимо самого графа должен быть доступен код для среды выполнения TensorFlow. Исполняемый код, созданный tfcompile , не использует среду выполнения TensorFlow и зависит только от ядер, которые фактически используются в вычислениях.
Компилятор построен на основе платформы XLA. Код, соединяющий TensorFlow с платформой XLA, находится в разделе tensorflow/compiler .
Что делает 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/compiler/aot/tests/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 "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;
}
Шаг 4. Создайте окончательный двоичный файл.
Этот шаг объединяет библиотеку, созданную tf_library на шаге 2, и код, написанный на шаге 3, для создания окончательного двоичного файла. Ниже приведен пример файла 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",
]
)