שימוש בקומפילציה של AOT

מה זה tfcompile?

tfcompile הוא כלי עצמאי כי קדימה-של-פעם (AOT) הידור TensorFlow גרפים לתוך קוד הפעלה. זה יכול להפחית את הגודל הבינארי הכולל, וגם למנוע כמה תקורה של זמן ריצה. טיפוסי לשימוש במקרה של tfcompile הוא לקמפל גרף היקש לתוך קוד הפעלה עבור התקנים ניידים.

גרף TensorFlow מבוצע בדרך כלל על ידי זמן הריצה של TensorFlow. זה כרוך בתקורת זמן ריצה מסוימת לביצוע של כל צומת בגרף. זה גם מוביל לגודל בינארי כולל גדול יותר, מכיוון שהקוד עבור זמן הריצה של TensorFlow צריך להיות זמין, בנוסף לגרף עצמו. קוד ההפעלה המיוצר על ידי tfcompile אינו משתמש ריצה TensorFlow, ורק יש תלות גרעינים המשמשים למעשה בחישוב.

המהדר בנוי על גבי מסגרת XLA. קוד גישור TensorFlow אל מתגורר במסגרת XLA תחת tensorflow / מהדר .

מה עושה tfcompile?

tfcompile לוקח subgraph, שזוהה על ידי מושגי TensorFlow של זנות ומביא, ומייצר פונקציה המיישם כי subgraph. feeds הן טיעוני הקלט עבור הפונקציה, ואת fetches הן טיעוני הפלט עבור הפונקציה. כל התשומות חייבות להיות מוגדרות במלואן על ידי העדכונים; תת-גרף הגזום שנוצר אינו יכול להכיל צמתים מציין מיקום או משתנה. מקובל לציין את כל מצייני המיקום והמשתנים כהזנות, מה שמבטיח שגרף המשנה המתקבל אינו מכיל יותר את הצמתים הללו. הפונקציה שנוצרה ארוז כמו cc_library , עם קובץ כותרת ייצוא חתימת הפונקציה, ואת קובץ אובייקט המכיל את היישום. המשתמש כותב קוד כדי להפעיל את הפונקציה שנוצרה לפי הצורך.

שימוש ב-tfcompile

סעיף זה מפרט את הפעולות ברמה גבוהה ליצירת בינארי הפעלה עם tfcompile מתוך subgraph TensorFlow. השלבים הם:

• שלב 1: הגדר את התת-גרף להידור
• שלב 2: השתמש tf_library מאקרו לבנות לקמפל את subgraph
• שלב 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 build מאקרו כדי להרכיב את תת-גרף

צעד זה ממיר את הגרף לתוך 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 לא יכול להרכיב subgraph המכילים Variables . Freeze_graph.py כלי הממיר משתנה לתוך קבוע, באמצעות ערכים מאוחסנים בקובץ במחסום. כתוצאה נוחות, tf_library המאקרו תומך freeze_checkpoint הטיעון, אשר מפעילה את הכלי. לדוגמאות נוספות ראה tensorflow / מהדר / AOT / בדיקות / לבנות .

קבועים המופיעים בתת-גרף הקומפילציה מורכבים ישירות לתוך הקוד שנוצר. כדי להעביר את הקבועים לפונקציה שנוצרה, במקום להרכיב אותם, פשוט העבירו אותם כהזנות.

לפרטים על 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 דומה, כשההבדל היחיד הוא השיטות לטיפול ב-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 לבנות קובץ.

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