روز جامعه ML 9 نوامبر است! برای به روز رسانی از TensorFlow، JAX به ما بپیوندید، و بیشتر بیشتر بدانید

با استفاده از تلفیق AOT

tfcompile چیست؟

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

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

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

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

tfcompile یک tfcompile می گیرد ، که توسط مفاهیم TensorFlow از فیدها و واکشی ها مشخص می شود و عملکردی را ایجاد می کند که آن زیرگراف را پیاده سازی می کند. feeds آرگومان های ورودی برای تابع و fetches ها آرگومان های خروجی برای تابع هستند. تمام ورودی ها باید به طور کامل توسط فیدها مشخص شوند. زیرشاخه هرس شده نمی تواند شامل گره های Placeholder یا Variable باشد. معمول است که همه متغیرهای متغیرهایی و متغیرها را به عنوان فید مشخص کنید ، این امر اطمینان حاصل می کند که فرعی حاصل دیگر حاوی این گره ها نیست. عملکرد تولید شده بصورت cc_library بسته بندی می شود ، یک فایل هدر ، امضای عملکرد را صادر می کند و یک پرونده شی حاوی پیاده سازی است. کاربر برای فراخوانی عملکرد ایجاد شده به طور مناسب کد را می نویسد.

با استفاده از tfcompile

این بخش جزئیات مراحل سطح بالا را برای تولید یک باینری قابل اجرا با tfcompile از tfcompile . مراحل عبارتند از:

  • مرحله 1: زیرنویس را برای کامپایل پیکربندی کنید
  • مرحله 2: از ماکرو ساخت 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 استفاده کنید

این مرحله با استفاده از ماکرو ساخت tf_library نمودار را به cc_library تبدیل می کند. cc_library شامل یک فایل شی object است که حاوی کد تولید شده از نمودار است ، همراه با یک فایل هدر که به کد تولید شده دسترسی دارد. tf_library با استفاده از tfcompile نمودار tfcompile را در کد اجرایی کامپایل می کند.

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 ) تولید شده توسط ماکرو build tf_library در مرحله قبل برای فراخوانی کد تولید شده استفاده می شود. پرونده هدر در فهرست bazel-bin مربوط به بسته build قرار دارد و براساس ویژگی 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 نمایانگر buffer های موقت که برای انجام محاسبات به صورت داخلی استفاده می شوند. به طور پیش فرض ، هر نمونه از کلاس تولید شده ، همه این بافرها را برای شما تخصیص داده و مدیریت می کند. برای تغییر این رفتار ممکن است از آرگومان سازنده 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",
    ]
)