Menggunakan kompilasi AOT

Apa itu tfcompile?

tfcompile adalah alat mandiri yang depan-of-waktu (AOT) mengkompilasi TensorFlow grafik ke dalam kode executable. Ini dapat mengurangi ukuran biner total, dan juga menghindari beberapa overhead runtime. Sebuah penggunaan-kasus yang khas dari tfcompile adalah untuk mengkompilasi sebuah grafik inferensi ke dalam kode executable untuk perangkat mobile.

Grafik TensorFlow biasanya dijalankan oleh runtime TensorFlow. Ini menimbulkan beberapa overhead runtime untuk eksekusi setiap node dalam grafik. Ini juga menghasilkan ukuran biner total yang lebih besar, karena kode untuk runtime TensorFlow harus tersedia, selain grafik itu sendiri. Kode dieksekusi diproduksi oleh tfcompile tidak menggunakan runtime TensorFlow, dan hanya memiliki ketergantungan pada kernel yang benar-benar digunakan dalam perhitungan.

Kompiler dibangun di atas kerangka XLA. Kode menjembatani TensorFlow ke XLA kerangka Resides bawah tensorflow / compiler .

Apa yang dilakukan tfcompile?

tfcompile mengambil subgraph, diidentifikasi oleh konsep TensorFlow feed dan fetch, dan menghasilkan fungsi yang mengimplementasikan bahwa subgraph. The feeds adalah argumen input untuk fungsi, dan fetches adalah argumen output untuk fungsi. Semua input harus sepenuhnya ditentukan oleh feed; subgraf yang dipangkas tidak dapat berisi node Placeholder atau Variabel. Adalah umum untuk menentukan semua Placeholder dan Variabel sebagai feed, yang memastikan subgraf yang dihasilkan tidak lagi berisi node ini. Fungsi yang dihasilkan dikemas sebagai cc_library , dengan file header mengekspor tanda tangan fungsi, dan file objek yang berisi pelaksanaannya. Pengguna menulis kode untuk memanggil fungsi yang dihasilkan sebagaimana mestinya.

Menggunakan tfcompile

Bagian ini menjelaskan langkah-langkah tingkat tinggi untuk menghasilkan executable biner dengan tfcompile dari subgraph TensorFlow. Langkah-langkahnya adalah:

  • Langkah 1: Konfigurasikan subgraf untuk dikompilasi
  • Langkah 2: Gunakan tf_library membangun makro untuk mengkompilasi subgraph yang
  • Langkah 3: Tulis kode untuk memanggil subgraf
  • Langkah 4: Buat biner terakhir

Langkah 1: Konfigurasikan subgraf untuk dikompilasi

Identifikasi umpan dan pengambilan yang sesuai dengan argumen input dan output untuk fungsi yang dihasilkan. Kemudian mengkonfigurasi feeds dan fetches dalam tensorflow.tf2xla.Config proto.

# 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" }
}

Langkah 2: Gunakan makro build tf_library untuk mengkompilasi subgraf

Langkah ini mengubah grafik menjadi cc_library menggunakan tf_library membangun makro. The cc_library terdiri dari file objek yang berisi kode yang dihasilkan dari grafik, bersama dengan file header yang memberikan akses ke kode yang dihasilkan. tf_library Memanfaatkan tfcompile untuk mengkompilasi grafik TensorFlow ke dalam kode executable.

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

Untuk menghasilkan GraphDef proto (test_graph_tfmatmul.pb) untuk contoh ini, jalankan make_test_graphs.py dan menentukan lokasi output dengan bendera --out_dir.

Grafik khas mengandung Variables yang mewakili bobot yang dipelajari melalui pelatihan, tetapi tfcompile tidak bisa mengkompilasi subgraf yang mengandung Variables . The freeze_graph.py alat mengkonversi variabel menjadi konstanta, menggunakan nilai yang disimpan dalam file pos pemeriksaan. Sebagai kenyamanan, yang tf_library makro mendukung freeze_checkpoint argumen, yang menjalankan alat. Untuk contoh lihat tensorflow / compiler / aot / tes / BUILD .

Konstanta yang muncul di subgraf yang dikompilasi dikompilasi langsung ke dalam kode yang dihasilkan. Untuk meneruskan konstanta ke fungsi yang dihasilkan, alih-alih mengkompilasinya, cukup berikan sebagai umpan.

Untuk rincian tentang tf_library membangun makro, lihat tfcompile.bzl .

Untuk rincian tentang mendasari tfcompile alat, lihat tfcompile_main.cc .

Langkah 3: Tulis kode untuk memanggil subgraf

Langkah ini menggunakan file header ( test_graph_tfmatmul.h ) yang dihasilkan oleh tf_library membangun makro pada langkah sebelumnya untuk memanggil kode yang dihasilkan. File header terletak di bazel-bin direktori yang sesuai dengan paket build, dan diberi nama berdasarkan nama atribut set pada tf_library membangun makro. Misalnya, header dihasilkan untuk test_graph_tfmatmul akan test_graph_tfmatmul.h . Di bawah ini adalah versi singkat dari apa yang dihasilkan. File yang dihasilkan, di bazel-bin , berisi komentar tambahan yang berguna.

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

Dihasilkan C ++ class disebut MatMulComp di foo::bar namespace, karena itu adalah cpp_class ditentukan dalam tf_library makro. Semua kelas yang dihasilkan memiliki API yang serupa, dengan satu-satunya perbedaan adalah metode untuk menangani buffer arg dan hasil. Metode tersebut berbeda berdasarkan jumlah dan jenis buffer, yang ditentukan oleh feed dan fetch argumen ke tf_library makro.

Ada tiga jenis buffer berhasil dalam kelas yang dihasilkan: args mewakili input, results mewakili output, dan temps mewakili buffer sementara yang digunakan secara internal untuk melakukan perhitungan. Secara default, setiap instance dari kelas yang dihasilkan mengalokasikan dan mengelola semua buffer ini untuk Anda. The AllocMode konstruktor argumen dapat digunakan untuk mengubah perilaku ini. Semua buffer disejajarkan dengan batas 64-byte.

Kelas C++ yang dihasilkan hanyalah pembungkus di sekitar kode tingkat rendah yang dihasilkan oleh XLA.

Contoh memohon fungsi yang dihasilkan berdasarkan 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;
}

Langkah 4: Buat biner terakhir

Langkah ini menggabungkan perpustakaan yang dihasilkan oleh tf_library pada langkah 2 dan kode yang ditulis dalam langkah 3 untuk membuat biner akhir. Berikut ini adalah contoh bazel berkas 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",
    ]
)