什么是tfcompile?
tfcompile
是一个独立的工具,提前(AOT)将TensorFlow图形编译成可执行代码。它可以减少总二进制大小,并且还可以避免一些运行时开销。典型的用例tfcompile
是将推理图编译成可移动设备的可执行代码。
TensorFlow图通常由TensorFlow运行时执行。这导致用于执行图中每个节点的一些运行时开销。这也导致更大的总二进制大小,因为除了图形本身之外,TensorFlow运行时的代码需要可用。所生成的可执行代码tfcompile
不使用TensorFlow运行时,并且只具有实际在计算中使用的内核的依赖关系。
编译器是建立在XLA框架之上的。将TensorFlow桥接到XLA框架的代码位于 tensorflow/compiler 下,它还包括对TensorFlow图形的即时(JIT)编译的支持。
tfcompile做什么?
tfcompile
采用由TensorFlow的饲料和提取概念识别的子图,并生成实现该子图的功能。这feeds
是函数的输入参数,它们是函数fetches
的输出参数。所有输入必须由供稿完全指定; 生成的修剪子图不能包含占位符或变量节点。通常将所有占位符和变量指定为Feed,这确保生成的子图不再包含这些节点。生成的函数打包成一个cc_library
头文件导出函数签名,以及一个包含该实现的目标文件。用户编写代码以适当地调用生成的函数。
使用tfcompile
本节详细介绍tfcompile
了使用TensorFlow子图生成可执行二进制文件的高级步骤 。步骤是:
- 步骤1:配置子图进行编译
- 步骤2:使用
tf_library
构建宏来编译子图 - 步骤3:编写代码以调用子图
- 步骤4:创建最终的二进制文件
步骤1:配置子图进行编译
识别与生成的函数的输入和输出参数相对应的馈送和提取。然后配置feeds
和fetches
在tensorflow.tfcompile.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("//third_party/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 proto(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-genfiles
与构建包相对应的目录中,并基于在tf_library
构建宏上设置的name属性命名。例如,生成的标题test_graph_tfmatmul
是test_graph_tfmatmul.h
。以下是生成的缩写版本。生成的文件bazel-genfiles
包含其他有用的注释。
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,唯一的区别是处理arg和结果缓冲区的方法。那些方法不同基于所述数量和类型的缓冲液,这是由所指定的feed
和fetch
参数给tf_library
宏。
在生成的类中有三种类型的缓冲区管理:args
表示输入,results
表示输出,并temps
表示在内部使用用于执行计算的临时缓冲区。默认情况下,生成的类的每个实例为您分配和管理所有这些缓冲区。该AllocMode
构造函数的参数可以用来改变这种行为。提供了一个便利库 tensorflow/compiler/aot/runtime.h
来帮助手动缓冲区分配; 此库的使用是可选的。所有缓冲区应与32字节边界对齐。
生成的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 //third_party/tensorflow/compiler/aot/tests/BUILD load("//third_party/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", ] )