יש שאלה? התחבר לקהילה בפורום הביקור של TensorFlow

שימוש באוסף AOT

מה זה tfcompile?

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 build tf_library כדי להרכיב את 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 באמצעות המאקרו לבנות 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 אינו יכול להרכיב תת תצלום המכיל Variables . הכלי freeze_graph.py ממיר משתנים לקבועים, תוך שימוש בערכים המאוחסנים בקובץ מחסום. tf_library המאקרו 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 ושמה מבוסס על מאפיין השם שהוגדר tf_library 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 נקראת MatMulComp במרחב השמות foo::bar , מכיוון שזה היה cpp_class שצוין במאקרו tf_library . לכל הכיתות שנוצרו יש ממשק API דומה, כאשר ההבדל היחיד הוא השיטות לטיפול בארגום ותאי מאגרים. אלה שיטות להשתנות בהתאם למספר סוגים של מאגרים, אשר פורטו על ידי feed ו fetch טיעונים אל tf_library מאקרו.

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

מחלקת C ++ שנוצרת היא רק עטיפה סביב הקוד ברמה נמוכה שנוצר על ידי XLA.

דוגמה tfcompile_test.cc הפונקציה שנוצרה על בסיס 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",
    ]
)