ใช้การคอมไพล์ AOT

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 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-bin ที่สอดคล้องกับแพ็คเกจ build และตั้งชื่อตามแอตทริบิวต์ชื่อที่ตั้งค่าไว้ในมาโคร build 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 ที่คล้ายกัน โดยมีความแตกต่างเพียงอย่างเดียวคือวิธีจัดการ arg และบัฟเฟอร์ผลลัพธ์ วิธีการเหล่านั้นแตกต่างกันไปตามจำนวนและประเภทของบัฟเฟอร์ ซึ่งระบุโดย 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",
    ]
)