דף זה תורגם על ידי Cloud Translation API.

שימוש באוסף תמונות AOT

מה זה tfcompile?

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

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

המהדר בנוי על גבי מסגרת XLA. הקוד שמגשר בין TensorFlow ל- XLA framework נמצא בקטע tensorflow/compiler.

מה עושה tfcompile?

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

שימוש ב-tfcompile

בקטע הזה מתוארים שלבים כלליים ליצירת קובץ בינארי של קובץ הפעלה עם tfcompile מתרשים משנה של TensorFlow. השלבים:

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

בשלב הזה, התרשים ממיר את התרשים ל-cc_library באמצעות המאקרו של build 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",
)

כדי ליצור את Proto של GraphDef (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.

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

פרטים נוספים על המאקרו של ה-build של tf_library מופיעים במאמר tfcompile.bzl.

פרטים על הכלי tfcompile שבבסיס הכלי הזה אפשר למצוא בכתובת tfcompile_main.cc.

שלב 3: כתיבת קוד להפעלת תת-הגרף

בשלב הזה, המערכת משתמשת בקובץ הכותרת (test_graph_tfmatmul.h) שנוצר על ידי המאקרו של ה-build של 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 דומה, כשההבדל היחיד הוא השיטות לטיפול במאגרים של ארגומנטים ותוצאות. השיטות האלה משתנות בהתאם למספר ולסוגים של מאגרי האחסון, שצוינו על ידי הארגומנטים feed ו-fetch במאקרו tf_library.

יש שלושה סוגי מאגרים שמנוהלים במחלקה שנוצרה: args לייצוג הקלט, results ייצוג של הפלטים ו-temps ייצוג של חוצצים זמניים שמשמשים באופן פנימי לביצוע החישוב. כברירת מחדל, כל מכונה של המחלקה שנוצרה מקצה ומנהלת את כל המאגרים האלה בשבילכם. אפשר להשתמש בארגומנט AllocMode constructor כדי לשנות את ההתנהגות. כל מאגרי הנתונים הזמניים מותאמים למגבלות של 64 בייטים.

מחלקת C++ שנוצרה היא רק wrapper סביב הקוד ברמה הנמוכה שנוצר על ידי 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 "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;
}

שלב 4: יוצרים את הקובץ הבינארי הסופי

השלב הזה משלב את הספרייה שנוצרה על ידי tf_library בשלב 2 עם הקוד שנכתב בשלב 3, כדי ליצור קובץ בינארי סופי. לפניכם דוגמה לקובץ BUILD מסוג 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",
    ]
)