استفاده از کامپایل AOT

tfcompile چیست؟

tfcompile یک ابزار مستقل است که پیش از زمان (AOT) نمودارهای TensorFlow را در کدهای اجرایی کامپایل می کند. این می تواند اندازه کل باینری را کاهش دهد و همچنین از برخی هزینه های سربار زمان اجرا جلوگیری کند. یک مورد معمول استفاده از tfcompile کامپایل یک نمودار استنتاج در کدهای اجرایی برای دستگاه های تلفن همراه است.

نمودار TensorFlow معمولاً توسط زمان اجرا TensorFlow اجرا می شود. این مقداری سربار زمان اجرا را برای اجرای هر گره در گراف متحمل می شود. این همچنین منجر به یک اندازه باینری کل بزرگتر می شود، زیرا کد مربوط به زمان اجرا TensorFlow، علاوه بر خود نمودار، باید در دسترس باشد. کد اجرایی تولید شده توسط tfcompile از زمان اجرا TensorFlow استفاده نمی کند و فقط به هسته هایی که واقعاً در محاسبات استفاده می شوند وابستگی دارد.

کامپایلر در بالای چارچوب XLA ساخته شده است. کد پل زدن TensorFlow به چارچوب XLA در زیر tensorflow/compiler قرار دارد.

tfcompile چه کاری انجام می دهد؟

tfcompile یک زیرگراف را می گیرد که با مفاهیم TensorFlow فیدها و واکشی ها شناسایی می شود و تابعی را ایجاد می کند که آن زیرگراف را پیاده سازی می کند. feeds آرگومان های ورودی برای تابع هستند و fetches آرگومان های خروجی برای تابع هستند. تمام ورودی ها باید به طور کامل توسط فیدها مشخص شوند. زیرگراف هرس شده نمی‌تواند حاوی گره‌های Placeholder یا Variable باشد. معمول است که همه مکان‌ها و متغیرها را به‌عنوان فید مشخص کنید، که تضمین می‌کند زیرگراف حاصل دیگر حاوی این گره‌ها نیست. تابع تولید شده به عنوان یک 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 برای کامپایل زیرگراف استفاده کنید

این مرحله با استفاده از ماکرو ساخت tf_library ، نمودار را به یک cc_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 از آرگومان freeze_checkpoint پشتیبانی می کند که ابزار را اجرا می کند. برای مثال‌های بیشتر به tensorflow/compiler/aot/tests/BUILD مراجعه کنید.

ثابت هایی که در زیرگراف کامپایل شده نشان داده می شوند مستقیماً در کد تولید شده کامپایل می شوند. برای ارسال ثابت ها به تابع تولید شده، به جای کامپایل کردن آنها، به سادگی آنها را به عنوان فید ارسال کنید.

برای جزئیات بیشتر در مورد ماکرو ساخت tf_library ، tfcompile.bzl را ببینید.

برای جزئیات بیشتر در مورد ابزار tfcompile زیربنایی، tfcompile_main.cc را ببینید.

مرحله 3: برای فراخوانی زیرگراف کد بنویسید

این مرحله از فایل هدر ( test_graph_tfmatmul.h ) تولید شده توسط ماکرو ساخت tf_library در مرحله قبل برای فراخوانی کد تولید شده استفاده می کند. فایل هدر در دایرکتوری bazel-bin مربوط به بسته ساخت قرار دارد و بر اساس ویژگی name مجموعه ای در ماکرو ساخت 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++ تولید شده در فضای نام foo::bar MatMulComp نامیده می شود، زیرا این 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 "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 را برای ایجاد یک باینری نهایی ترکیب می کند. در زیر یک نمونه فایل 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",
    ]
)