Google I/O is a wrap! Catch up on TensorFlow sessions View sessions

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

Tfcompile là gì?

tfcompile là một công cụ độc lập phía trước-of-time (AOT) biên dịch TensorFlow đồ thị vào mã thực thi. Nó có thể làm giảm tổng kích thước nhị phân và cũng tránh một số chi phí thời gian chạy. Một use-case điển hình của tfcompile là để biên dịch một đồ thị suy luận thành mã thực thi cho các 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 phát sinh một số chi phí thời gian chạy để thực hiện từng nút trong biểu đồ. Điều này cũng dẫn đến tổng kích thước nhị phân lớn hơn, vì mã cho thời gian chạy TensorFlow cần phải có sẵn, ngoài bản thân biểu đồ. Các mã thực thi được tạo ra bởi tfcompile không sử dụng thời gian chạy TensorFlow, và chỉ có phụ thuộc vào hạt nhân đó đang thực sự được sử dụng trong việc tính toán.

Trình biên dịch được xây dựng dựa trên khung XLA. Mã lấp TensorFlow đến cư trú khuôn khổ XLA dưới tensorflow / biên dịch .

Tfcompile làm gì?

tfcompile mất một đồ thị con, xác định bởi các khái niệm TensorFlow thức ăn chăn nuôi và fetches, và tạo ra một chức năng mà cụ rằng đồ thị con. Các feeds là các đối số đầu vào cho các chức năng, và các fetches là các đối số đầu ra cho các chức năng. Tất cả các đầu vào phải được xác định đầy đủ bởi các nguồn cấp dữ liệu; đồ thị con được lược bớt kết quả không được chứa các nút Giữ chỗ hoặc Biến. Thông thường, chỉ định tất cả Trình giữ chỗ và Biến làm nguồn cấp dữ liệu, điều này đảm bảo rằng đồ thị con kết quả không còn chứa các nút này nữa. Chức năng tạo được đóng gói như một cc_library , với một tập tin header xuất khẩu chữ ký chức năng, và một đối tượng tập tin chứa thi. Người dùng viết mã để gọi hàm được tạo nếu thích hợp.

Sử dụng tfcompile

Phần này chi tiết các bước trình độ cao để tạo ra một thực thi nhị phân với tfcompile từ một đồ thị con TensorFlow. Các bước là:

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

Bước 1: Định cấu hình biểu đồ 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 đó cấu hình các feedsfetches trong một 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" }
}

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

Bước này chuyển đổi đồ thị thành một cc_library sử dụng tf_library xây dựng vĩ mô. Các cc_library bao gồm một tập tin đối tượng có chứa các mã được tạo từ biểu đồ, cùng với một tập tin tiêu đề mà cho phép truy cập vào mã được tạo ra. tf_library sử dụng tfcompile để biên dịch đồ thị TensorFlow vào mã 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) ví dụ này, chạy make_test_graphs.py và xác định vị trí đầu ra với cờ --out_dir.

Đồ thị điển hình chứa Variables đại diện cho trọng lượng được học qua đào tạo, nhưng tfcompile không thể biên dịch một đồ thị con có chứa Variables . Các freeze_graph.py công cụ chuyển đổi biến thành hằng số, sử dụng các giá trị được lưu trữ trong một tập tin trạm kiểm soát. Để thuận tiện, các tf_library vĩ mô hỗ trợ freeze_checkpoint luận, chạy công cụ. Để biết thêm ví dụ thấy tensorflow / biên dịch / AOT / kiểm tra / BUILD .

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

Để biết chi tiết về các tf_library xây dựng vĩ mô, xem tfcompile.bzl .

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

Bước 3: Viết mã để gọi ra đoạn con

Bước này sử dụng các tập tin tiêu đề ( test_graph_tfmatmul.h ) tạo ra bởi các tf_library xây dựng vĩ mô trong bước trước để gọi mã được tạo. Các tập tin tiêu đề nằm ở bazel-bin thư mục tương ứng với việc xây dựng gói, và được đặt tên dựa trên các thiết lập tên thuộc tính trên tf_library xây dựng vĩ mô. Ví dụ, tiêu đề được tạo ra cho test_graph_tfmatmul sẽ test_graph_tfmatmul.h . Dưới đây là phiên bản viết tắt của những gì được tạo. Các tập tin được tạo ra, trong bazel-bin , chứa ý kiến bổ sung 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

Được tạo ra lớp C ++ được gọi MatMulComp trong foo::bar namespace, bởi vì đó là cpp_class quy định tại các tf_library vĩ mô. Tất cả các lớp được tạo đều có một API giống nhau, với sự khác biệt duy nhất là các phương thức để xử lý bộ đệm arg và kết quả. Những phương pháp khác nhau dựa trên số lượng và loại của bộ đệm, được xác định bởi các feedfetch đối số cho các tf_library vĩ mô.

Có ba loại bộ đệm quản lý trong lớp được tạo ra: args đại diện cho đầu vào, results đại diện cho kết quả đầu ra, và temps đại diện cho bộ đệm tạm thời sử dụng trong nội bộ để thực hiện các tính toán. Theo mặc định, mỗi phiên bản 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. Các AllocMode luận constructor có thể được sử dụng để thay đổi hành vi này. Tất cả các bộ đệm được căn chỉnh theo ranh giới 64 byte.

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

Ví dụ về cách gọi các chức năng được tạo ra 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 "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 các thư viện được tạo ra bởi tf_library trong bước 2 và mã được viết ở bước 3 để tạo ra một nhị phân thức. Dưới đây là một ví dụ bazel tập tin 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",
    ]
)