Sử dụng biên dịch AOT

tfcompile là gì?

tfcompile là một công cụ độc lập có chức năng biên dịch trước (AOT) các biểu đồ TensorFlow thành mã có thể thực thi. Việc này có thể làm giảm tổng kích thước tệp nhị phân và cũng tránh được một số hao tổn thời gian chạy. Một trường hợp sử dụng điển hình của tfcompile là biên dịch biểu đồ suy luận thành mã có thể thực thi cho thiết bị di động.

Biểu đồ TensorFlow thường được thực thi bởi thời gian chạy TensorFlow. Điều này làm phát sinh một số hao tổn thời gian chạy khi thực thi từng nút trong biểu đồ. Điều này cũng dẫn đến tổng kích thước tệp nhị phân lớn hơn, vì ngoài chính biểu đồ, TensorFlow cần có sẵn mã cho thời gian chạy TensorFlow. Mã thực thi do tfcompile tạo ra không sử dụng môi trường thời gian chạy TensorFlow và chỉ có các phần phụ thuộc trên các hạt nhân thực sự được dùng trong quá trình tính toán.

Trình biên dịch được xây dựng dựa trên khung XLA. Đoạn mã cầu nối TensorFlow đến khung XLA nằm trong tensorflow/compiler.

Tfcompile có chức năng gì?

tfcompile lấy một đồ thị con, được xác định theo các khái niệm của TensorFlow về nguồn cấp dữ liệu và tìm nạp, rồi tạo một hàm triển khai đồ thị con đó. feeds là các đối số đầu vào cho hàm và fetches là các đối số đầu ra cho hàm. Tất cả dữ liệu đầu vào phải được nguồn cấp dữ liệu chỉ định đầy đủ; biểu đồ con bị cắt bớt không thể chứa các nút Phần giữ chỗ hoặc Biến. Thường thì bạn nên chỉ định tất cả các Trình giữ chỗ và Biến làm nguồn cấp dữ liệu. Điều này đảm bảo đồ thị con thu được không còn chứa các nút này nữa. Hàm đã tạo được đóng gói dưới dạng cc_library, với một tệp tiêu đề xuất chữ ký hàm và một tệp đối tượng chứa phương thức triển khai. Người dùng viết mã để gọi hàm đã tạo khi phù hợp.

Sử dụng tfcompile

Phần này trình bày chi tiết các bước cấp cao để tạo tệp nhị phân có thể thực thi bằng tfcompile từ biểu đồ con TensorFlow. Các bước thực hiện:

• Bước 1: Định cấu hình đồ thị con để biên dịch
• Bước 2: Sử dụng macro bản dựng tf_library để biên dịch đồ thị con
• Bước 3: Viết mã để gọi đồ thị con
• Bước 4: Tạo tệp nhị phân cuối cùng

Bước 1: Định cấu hình đồ thị con để biên dịch

Xác định các nguồn cấp dữ liệu và tìm nạp tương ứng với các đối số đầu vào và đầu ra cho hàm đã tạo. Sau đó, hãy định cấu hình feeds và fetches trong một proto 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" }
}

Bước 2: Sử dụng macro dựng tf_library để biên dịch đồ thị con

Bước này chuyển đổi biểu đồ thành cc_library bằng cách sử dụng macro bản dựng tf_library. cc_library bao gồm một tệp đối tượng chứa mã được tạo từ biểu đồ, cùng với một tệp tiêu đề cấp quyền truy cập vào mã đã tạo. tf_library sử dụng tfcompile để biên dịch biểu đồ TensorFlow thành mã có thể thực thi.

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

Để tạo proto GraphDef (test_graph_tfmatmul.pb) cho ví dụ này, hãy chạy make_test_graphs.py và chỉ định vị trí đầu ra bằng cờ --out_dir.

Các biểu đồ thông thường chứa Variables đại diện cho trọng số học được thông qua quá trình huấn luyện, nhưng tfcompile không thể biên dịch đồ thị con chứa Variables. Công cụ freeze_graph.py chuyển đổi các biến thành hằng số, sử dụng các giá trị được lưu trữ trong tệp điểm kiểm tra. Để thuận tiện, macro tf_library hỗ trợ đối số freeze_checkpoint chạy công cụ. Để biết thêm ví dụ, hãy xem bài viết tensorflow/compiler/aot/tests/BUILD.

Các hằng số hiện trong biểu đồ con đã biên dịch được biên dịch trực tiếp vào mã đã tạo. Để truyền các hằng số vào hàm đã tạo, thay vì để chúng được biên dịch, bạn chỉ cần truyền các hằng số đó vào dưới dạng nguồn cấp dữ liệu.

Để biết thông tin chi tiết về macro bản dựng tf_library, vui lòng xem tfcompile.bzl.

Để biết thông tin chi tiết về công cụ tfcompile cơ bản, hãy xem tfcompile_main.cc.

Bước 3: Viết mã để gọi đồ thị con

Bước này sử dụng tệp tiêu đề (test_graph_tfmatmul.h) do macro bản dựng tf_library tạo trong bước trước để gọi mã đã tạo. Tệp tiêu đề nằm trong thư mục bazel-bin tương ứng với gói bản dựng và được đặt tên dựa trên thuộc tính tên được đặt trong macro bản dựng tf_library. Ví dụ: tiêu đề được tạo cho test_graph_tfmatmul sẽ là test_graph_tfmatmul.h. Dưới đây là phiên bản viết tắt của những dữ liệu được tạo. Tệp được tạo trong bazel-bin có chứa thêm các nhận xét hữu ích.

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

Lớp C++ đã tạo có tên là MatMulComp trong không gian tên foo::bar, vì đó là cpp_class được chỉ định trong macro tf_library. Tất cả các lớp được tạo đều có API tương tự, chỉ khác là các phương thức để xử lý vùng đệm đối số và kết quả. Các phương thức đó khác nhau tuỳ theo số lượng và loại vùng đệm, được chỉ định bằng các đối số feed và fetch cho macro tf_library.

Có 3 loại vùng đệm được quản lý trong lớp được tạo: args đại diện cho đầu vào, results đại diện cho đầu ra và temps đại diện cho vùng đệm tạm thời được dùng nội bộ để thực hiện việc tính toán. Theo mặc định, mỗi thực thể của lớp được tạo sẽ phân bổ và quản lý tất cả các vùng đệm này cho bạn. Bạn có thể dùng đối số hàm khởi tạo AllocMode để thay đổi hành vi này. Tất cả các vùng đệm đều được căn chỉnh theo ranh giới 64 byte.

Lớp C++ được tạo chỉ là trình bao bọc xung quanh mã cấp thấp do XLA tạo.

Ví dụ về cách gọi hàm được tạo dựa trên 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;
}

Bước 4: Tạo tệp nhị phân cuối cùng

Bước này kết hợp thư viện do tf_library tạo ở bước 2 và mã được viết ở bước 3 để tạo tệp nhị phân cuối cùng. Dưới đây là tệp bazel XÂY DỰNG mẫu.

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